{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![QuantConnect Logo](https://cdn.quantconnect.com/web/i/icon.png)\n",
    "<hr>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def data_fetch(start_time, end_time, stock_list, resolution=Resolution.Daily) :\n",
    "    #resolution: Daily, Hour, Minute, Second\n",
    "    qb = QuantBook()\n",
    "    for stock in stock_list:\n",
    "        qb.AddEquity(stock)\n",
    "    history = qb.History(qb.Securities.Keys, start_time, end_time, resolution)\n",
    "    return history\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "def preprocess(df, stock_list):\n",
    "    df = df[['open','high','low','close','volume']]\n",
    "    if_first_time = True\n",
    "    for stock in stock_list:\n",
    "        if if_first_time:\n",
    "            ary = df.loc[stock].values\n",
    "            if_first_time = False\n",
    "        else:\n",
    "            temp = df.loc[stock].values\n",
    "            ary = np.hstack((ary,temp))\n",
    "    return ary\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "class StockTradingEnv():\n",
    "    \"\"\"FinRL\n",
    "    Paper: A Deep Reinforcement Learning Library for Automated Stock Trading in Quantitative Finance\n",
    "           https://arxiv.org/abs/2011.09607 NeurIPS 2020: Deep RL Workshop.\n",
    "    Source: Github https://github.com/AI4Finance-LLC/FinRL-Library\n",
    "    Modify: Github Yonv1943 ElegantRL\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, ary, initial_account=1e6, max_stock=1e2, transaction_fee_percent=1e-3, if_train=True,\n",
    "                 ):\n",
    "        self.stock_dim = 3\n",
    "        self.initial_account = initial_account\n",
    "        self.transaction_fee_percent = transaction_fee_percent\n",
    "        self.max_stock = max_stock\n",
    "\n",
    "        N = ary.shape[0]\n",
    "        x = int(N*0.8)# ary: (date, item*stock_dim), item: (adjcp, macd, rsi, cci, adx)\n",
    "        \n",
    "        # reset\n",
    "        self.day = 0\n",
    "        if if_train:\n",
    "            self.ary = ary[:x] \n",
    "        else:\n",
    "            self.ary = ary[x:]\n",
    "        self.initial_account__reset = self.initial_account\n",
    "        self.account = self.initial_account__reset\n",
    "        self.day_npy = self.ary[self.day]\n",
    "        self.prices = self.day_npy[[5*x+3 for x in range(self.stock_dim)]]\n",
    "        self.stocks = np.zeros(self.stock_dim, dtype=np.float32)  # multi-stack\n",
    "\n",
    "        self.total_asset = self.account + (self.prices[:self.stock_dim] * self.stocks).sum()\n",
    "        self.episode_return = 0.0  # Compatibility for ElegantRL 2020-12-21\n",
    "        self.gamma_return = 0.0\n",
    "\n",
    "        '''env information'''\n",
    "        self.env_name = 'Stock_qc-v1'\n",
    "        self.state_dim = 1 + 6 * self.stock_dim\n",
    "        self.action_dim = self.stock_dim\n",
    "        self.if_discrete = False\n",
    "        self.target_return = 1.25  # convergence 1.5\n",
    "        self.max_step = self.ary.shape[0]\n",
    "\n",
    "    def reset(self):\n",
    "        self.account = self.initial_account\n",
    "        self.stocks = np.zeros(self.stock_dim, dtype=np.float32)\n",
    "        self.prices = self.day_npy[[5*x+3 for x in range(self.stock_dim)]]\n",
    "        self.total_asset = self.account + (self.prices[:self.stock_dim] * self.stocks).sum()\n",
    "        # total_asset = account + (adjcp * stocks).sum()\n",
    "\n",
    "        self.day = 0\n",
    "        self.day_npy = self.ary[self.day]\n",
    "        self.day += 1\n",
    "\n",
    "        state = np.hstack((self.account * 2 ** -16,\n",
    "                           self.day_npy * 2 ** -8,\n",
    "                           self.stocks * 2 ** -12,), ).astype(np.float32)\n",
    "        return state\n",
    "\n",
    "    def step(self, action):\n",
    "        action = action * self.max_stock\n",
    "        self.prices = self.day_npy[[5*x+3 for x in range(self.stock_dim)]]\n",
    "        \"\"\"buy or sell stock\"\"\"\n",
    "        for index in range(self.stock_dim):\n",
    "            stock_action = action[index]\n",
    "            adj = self.prices[index]\n",
    "            if stock_action > 0:  # buy_stock\n",
    "                available_amount = self.account // adj\n",
    "                delta_stock = min(available_amount, stock_action)\n",
    "                self.account -= adj * delta_stock * (1 + self.transaction_fee_percent)\n",
    "                self.stocks[index] += delta_stock\n",
    "            elif self.stocks[index] > 0:  # sell_stock\n",
    "                delta_stock = min(-stock_action, self.stocks[index])\n",
    "                self.account += adj * delta_stock * (1 - self.transaction_fee_percent)\n",
    "                self.stocks[index] -= delta_stock\n",
    "\n",
    "        \"\"\"update day\"\"\"\n",
    "        self.day_npy = self.ary[self.day]\n",
    "        self.prices = self.day_npy[[5*x+3 for x in range(self.stock_dim)]]\n",
    "        self.day += 1\n",
    "        done = self.day == self.max_step  # 2020-12-21\n",
    "\n",
    "        state = np.hstack((self.account * 2 ** -16,\n",
    "                           self.day_npy * 2 ** -8,\n",
    "                           self.stocks * 2 ** -12,), ).astype(np.float32)\n",
    "\n",
    "        next_total_asset = self.account + (self.prices[:self.stock_dim] * self.stocks).sum()\n",
    "        reward = (next_total_asset - self.total_asset) * 2 ** -16  # notice scaling!\n",
    "        self.total_asset = next_total_asset\n",
    "\n",
    "        self.gamma_return = self.gamma_return * 0.99 + reward  # notice: gamma_r seems good? Yes\n",
    "        if done:\n",
    "            reward += self.gamma_return\n",
    "            self.gamma_return = 0.0  # env.reset()\n",
    "        \n",
    "            # cumulative_return_rate\n",
    "            self.episode_return = next_total_asset / self.initial_account\n",
    "            print(self.episode_return)\n",
    "        return state, reward, done, None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "\n",
    "\"\"\"[ElegantRL](https://github.com/AI4Finance-LLC/ElegantRL)\"\"\"\n",
    "\n",
    "'''Q Network'''\n",
    "\n",
    "class Hardswish(nn.Module):  # export-friendly version of nn.Hardswish()\n",
    "    @staticmethod\n",
    "    def forward(x):\n",
    "        # return x * F.hardsigmoid(x)  # for torchscript and CoreML\n",
    "        return x * nn.functional.hardtanh(x + 3, 0., 6.) / 6.  # for torchscript, CoreML and ONNX\n",
    "\n",
    "        \n",
    "class QNet(nn.Module):  # nn.Module is a standard PyTorch Network\n",
    "    def __init__(self, mid_dim, state_dim, action_dim):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(nn.Linear(state_dim, mid_dim), nn.ReLU(),\n",
    "                                 nn.Linear(mid_dim, mid_dim), nn.ReLU(),\n",
    "                                 nn.Linear(mid_dim, mid_dim), nn.ReLU(),\n",
    "                                 nn.Linear(mid_dim, action_dim))\n",
    "\n",
    "    def forward(self, state):\n",
    "        return self.net(state)  # Q value\n",
    "\n",
    "\n",
    "class QNetDuel(nn.Module):  # Dueling DQN\n",
    "    def __init__(self, mid_dim, state_dim, action_dim):\n",
    "        super().__init__()\n",
    "        self.net_state = nn.Sequential(nn.Linear(state_dim, mid_dim), nn.ReLU(),\n",
    "                                       nn.Linear(mid_dim, mid_dim), nn.ReLU())\n",
    "        self.net_val = nn.Sequential(nn.Linear(mid_dim, mid_dim), nn.ReLU(),\n",
    "                                     nn.Linear(mid_dim, 1))  # Q value\n",
    "        self.net_adv = nn.Sequential(nn.Linear(mid_dim, mid_dim), nn.ReLU(),\n",
    "                                     nn.Linear(mid_dim, action_dim))  # advantage function value 1\n",
    "\n",
    "    def forward(self, state):\n",
    "        t_tmp = self.net_state(state)\n",
    "        q_val = self.net_val(t_tmp)\n",
    "        q_adv = self.net_adv(t_tmp)\n",
    "        return q_val + q_adv - q_adv.mean(dim=1, keepdim=True)  # dueling Q value\n",
    "\n",
    "\n",
    "class QNetTwin(nn.Module):  # Double DQN\n",
    "    def __init__(self, mid_dim, state_dim, action_dim):\n",
    "        super().__init__()\n",
    "        self.net_state = nn.Sequential(nn.Linear(state_dim, mid_dim), nn.ReLU(),\n",
    "                                       nn.Linear(mid_dim, mid_dim), nn.ReLU())  # state\n",
    "        self.net_q1 = nn.Sequential(nn.Linear(mid_dim, mid_dim), nn.ReLU(),\n",
    "                                    nn.Linear(mid_dim, action_dim))  # q1 value\n",
    "        self.net_q2 = nn.Sequential(nn.Linear(mid_dim, mid_dim), nn.ReLU(),\n",
    "                                    nn.Linear(mid_dim, action_dim))  # q2 value\n",
    "\n",
    "    def forward(self, state):\n",
    "        tmp = self.net_state(state)\n",
    "        return self.net_q1(tmp)  # one Q value\n",
    "\n",
    "    def get_q1_q2(self, state):\n",
    "        tmp = self.net_state(state)\n",
    "        q1 = self.net_q1(tmp)\n",
    "        q2 = self.net_q2(tmp)\n",
    "        return q1, q2  # two Q values\n",
    "\n",
    "\n",
    "class QNetTwinDuel(nn.Module):  # D3QN: Dueling Double DQN\n",
    "    def __init__(self, mid_dim, state_dim, action_dim):\n",
    "        super().__init__()\n",
    "        self.net_state = nn.Sequential(nn.Linear(state_dim, mid_dim), nn.ReLU(),\n",
    "                                       nn.Linear(mid_dim, mid_dim), nn.ReLU())\n",
    "        self.net_val1 = nn.Sequential(nn.Linear(mid_dim, mid_dim), nn.ReLU(),\n",
    "                                      nn.Linear(mid_dim, 1))  # q1 value\n",
    "        self.net_val2 = nn.Sequential(nn.Linear(mid_dim, mid_dim), nn.ReLU(),\n",
    "                                      nn.Linear(mid_dim, 1))  # q2 value\n",
    "        self.net_adv1 = nn.Sequential(nn.Linear(mid_dim, mid_dim), nn.ReLU(),\n",
    "                                      nn.Linear(mid_dim, action_dim))  # advantage function value 1\n",
    "        self.net_adv2 = nn.Sequential(nn.Linear(mid_dim, mid_dim), nn.ReLU(),\n",
    "                                      nn.Linear(mid_dim, action_dim))  # advantage function value 1\n",
    "\n",
    "    def forward(self, state):\n",
    "        t_tmp = self.net_state(state)\n",
    "        q_val = self.net_val1(t_tmp)\n",
    "        q_adv = self.net_adv1(t_tmp)\n",
    "        return q_val + q_adv - q_adv.mean(dim=1, keepdim=True)  # one dueling Q value\n",
    "\n",
    "    def get_q1_q2(self, state):\n",
    "        tmp = self.net_state(state)\n",
    "\n",
    "        val1 = self.net_val1(tmp)\n",
    "        adv1 = self.net_adv1(tmp)\n",
    "        q1 = val1 + adv1 - adv1.mean(dim=1, keepdim=True)\n",
    "\n",
    "        val2 = self.net_val2(tmp)\n",
    "        adv2 = self.net_adv2(tmp)\n",
    "        q2 = val2 + adv2 - adv2.mean(dim=1, keepdim=True)\n",
    "        return q1, q2  # two dueling Q values\n",
    "\n",
    "\n",
    "'''Policy Network (Actor)'''\n",
    "\n",
    "\n",
    "class Actor(nn.Module):  # DPG: Deterministic Policy Gradient\n",
    "    def __init__(self, mid_dim, state_dim, action_dim):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(nn.Linear(state_dim, mid_dim), nn.ReLU(),\n",
    "                                 nn.Linear(mid_dim, mid_dim), nn.ReLU(),\n",
    "                                 nn.Linear(mid_dim, mid_dim), nn.ReLU(),\n",
    "                                 nn.Linear(mid_dim, action_dim))\n",
    "\n",
    "    def forward(self, state):\n",
    "        return self.net(state).tanh()  # action.tanh()\n",
    "\n",
    "    def get_action(self, state, action_std):\n",
    "        action = self.net(state).tanh()\n",
    "        noise = (torch.randn_like(action) * action_std).clamp(-0.5, 0.5)\n",
    "        return (action + noise).clamp(-1.0, 1.0)\n",
    "\n",
    "\n",
    "class ActorPPO(nn.Module):\n",
    "    def __init__(self, mid_dim, state_dim, action_dim, if_use_dn=False):\n",
    "        super().__init__()\n",
    "        if isinstance(state_dim, int):\n",
    "            if if_use_dn:\n",
    "                nn_dense = DenseNet(mid_dim // 2)\n",
    "                inp_dim = nn_dense.inp_dim\n",
    "                out_dim = nn_dense.out_dim\n",
    "\n",
    "                self.net = nn.Sequential(nn.Linear(state_dim, inp_dim), nn.ReLU(),\n",
    "                                         nn_dense,\n",
    "                                         nn.Linear(out_dim, action_dim), )\n",
    "            else:\n",
    "                self.net = nn.Sequential(nn.Linear(state_dim, mid_dim), nn.ReLU(),\n",
    "                                         nn.Linear(mid_dim, mid_dim), nn.ReLU(),\n",
    "                                         nn.Linear(mid_dim, mid_dim), Hardswish(),\n",
    "                                         nn.Linear(mid_dim, action_dim), )\n",
    "        else:\n",
    "            def set_dim(i):\n",
    "                return int(12 * 1.5 ** i)\n",
    "\n",
    "            self.net = nn.Sequential(NnReshape(*state_dim),  # -> [batch_size, 4, 96, 96]\n",
    "                                     nn.Conv2d(state_dim[0], set_dim(0), 4, 2, bias=True), nn.LeakyReLU(),\n",
    "                                     nn.Conv2d(set_dim(0), set_dim(1), 3, 2, bias=False), nn.ReLU(),\n",
    "                                     nn.Conv2d(set_dim(1), set_dim(2), 3, 2, bias=False), nn.ReLU(),\n",
    "                                     nn.Conv2d(set_dim(2), set_dim(3), 3, 2, bias=True), nn.ReLU(),\n",
    "                                     nn.Conv2d(set_dim(3), set_dim(4), 3, 1, bias=True), nn.ReLU(),\n",
    "                                     nn.Conv2d(set_dim(4), set_dim(5), 3, 1, bias=True), nn.ReLU(),\n",
    "                                     NnReshape(-1),\n",
    "                                     nn.Linear(set_dim(5), mid_dim), nn.ReLU(),\n",
    "                                     nn.Linear(mid_dim, action_dim), )\n",
    "\n",
    "        self.a_std_log = nn.Parameter(torch.zeros((1, action_dim)) - 0.5, requires_grad=True)  # trainable parameter\n",
    "        self.sqrt_2pi_log = np.log(np.sqrt(2 * np.pi))\n",
    "\n",
    "        layer_norm(self.net[-1], std=0.1)  # output layer for action\n",
    "\n",
    "    def forward(self, state):\n",
    "        return self.net(state).tanh()  # action\n",
    "\n",
    "    def get_action_noise(self, state):\n",
    "        a_avg = self.net(state)\n",
    "        a_std = self.a_std_log.exp()\n",
    "\n",
    "        noise = torch.randn_like(a_avg)\n",
    "        action = a_avg + noise * a_std\n",
    "        return action, noise\n",
    "\n",
    "    def compute_logprob(self, state, action):\n",
    "        a_avg = self.net(state)\n",
    "        a_std = self.a_std_log.exp()\n",
    "        delta = ((a_avg - action) / a_std).pow(2).__mul__(0.5)  # __mul__(0.5) is * 0.5\n",
    "        logprob = -(self.a_std_log + self.sqrt_2pi_log + delta)\n",
    "        return logprob.sum(1)\n",
    "\n",
    "\n",
    "class ActorSAC(nn.Module):\n",
    "    def __init__(self, mid_dim, state_dim, action_dim, if_use_dn=False):\n",
    "        super().__init__()\n",
    "        if if_use_dn:\n",
    "            nn_dense = DenseNet(mid_dim // 2)\n",
    "            inp_dim = nn_dense.inp_dim\n",
    "            out_dim = nn_dense.out_dim\n",
    "\n",
    "            self.net_state = nn.Sequential(nn.Linear(state_dim, inp_dim), nn.ReLU(),\n",
    "                                           nn_dense, )\n",
    "        else:\n",
    "            self.net_state = nn.Sequential(nn.Linear(state_dim, mid_dim), nn.ReLU(),\n",
    "                                           nn.Linear(mid_dim, mid_dim), Hardswish(),\n",
    "                                           nn.Linear(mid_dim, mid_dim), Hardswish())\n",
    "            out_dim = mid_dim\n",
    "\n",
    "        self.net_a_avg = nn.Linear(out_dim, action_dim)  # the average of action\n",
    "        self.net_a_std = nn.Linear(out_dim, action_dim)  # the log_std of action\n",
    "\n",
    "        self.sqrt_2pi_log = np.log(np.sqrt(2 * np.pi))\n",
    "        layer_norm(self.net_a_avg, std=0.01)  # output layer for action, it is no necessary.\n",
    "\n",
    "    def forward(self, state):\n",
    "        tmp = self.net_state(state)\n",
    "        return self.net_a_avg(tmp).tanh()  # action\n",
    "\n",
    "    def get_action(self, state):\n",
    "        t_tmp = self.net_state(state)\n",
    "        a_avg = self.net_a_avg(t_tmp)  # NOTICE! it is a_avg without .tanh()\n",
    "        a_std = self.net_a_std(t_tmp).clamp(-20, 2).exp()\n",
    "        return torch.normal(a_avg, a_std).tanh()  # re-parameterize\n",
    "\n",
    "    def get_action_logprob(self, state):\n",
    "        t_tmp = self.net_state(state)\n",
    "        a_avg = self.net_a_avg(t_tmp)  # NOTICE! it needs a_avg.tanh()\n",
    "        a_std_log = self.net_a_std(t_tmp).clamp(-20, 2)\n",
    "        a_std = a_std_log.exp()\n",
    "\n",
    "        \"\"\"add noise to action in stochastic policy\"\"\"\n",
    "        noise = torch.randn_like(a_avg, requires_grad=True)\n",
    "        action = a_avg + a_std * noise\n",
    "        a_tan = action.tanh()  # action.tanh()\n",
    "        # Can only use above code instead of below, because the tensor need gradients here.\n",
    "        # a_noise = torch.normal(a_avg, a_std, requires_grad=True)\n",
    "\n",
    "        '''compute logprob according to mean and std of action (stochastic policy)'''\n",
    "        # # self.sqrt_2pi_log = np.log(np.sqrt(2 * np.pi))\n",
    "        # logprob = a_std_log + self.sqrt_2pi_log + noise.pow(2).__mul__(0.5)  # noise.pow(2) * 0.5\n",
    "        # different from above (gradient)\n",
    "        delta = ((a_avg - action) / a_std).pow(2).__mul__(0.5)\n",
    "        logprob = a_std_log + self.sqrt_2pi_log + delta\n",
    "        # same as below:\n",
    "        # from torch.distributions.normal import Normal\n",
    "        # logprob_noise = Normal(a_avg, a_std).logprob(a_noise)\n",
    "        # logprob = logprob_noise + (-a_noise_tanh.pow(2) + 1.000001).log()\n",
    "        # same as below:\n",
    "        # a_delta = (a_avg - a_noise).pow(2) /(2*a_std.pow(2))\n",
    "        # logprob_noise = -a_delta - a_std.log() - np.log(np.sqrt(2 * np.pi))\n",
    "        # logprob = logprob_noise + (-a_noise_tanh.pow(2) + 1.000001).log()\n",
    "\n",
    "        logprob = logprob + (-a_tan.pow(2) + 1.000001).log()  # fix logprob using the derivative of action.tanh()\n",
    "        # same as below:\n",
    "        # epsilon = 1e-6\n",
    "        # logprob = logprob_noise - (1 - a_noise_tanh.pow(2) + epsilon).log()\n",
    "        return a_tan, logprob.sum(1, keepdim=True)\n",
    "\n",
    "\n",
    "'''Value Network (Critic)'''\n",
    "\n",
    "\n",
    "class Critic(nn.Module):\n",
    "    def __init__(self, mid_dim, state_dim, action_dim):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(nn.Linear(state_dim + action_dim, mid_dim), nn.ReLU(),\n",
    "                                 nn.Linear(mid_dim, mid_dim), nn.ReLU(),\n",
    "                                 nn.Linear(mid_dim, 1))\n",
    "\n",
    "    def forward(self, state, action):\n",
    "        return self.net(torch.cat((state, action), dim=1))  # Q value\n",
    "\n",
    "\n",
    "class CriticAdv(nn.Module):\n",
    "    def __init__(self, state_dim, mid_dim, if_use_dn=False):\n",
    "        super().__init__()\n",
    "        if isinstance(state_dim, int):\n",
    "            if if_use_dn:\n",
    "                nn_dense = DenseNet(mid_dim // 2)\n",
    "                inp_dim = nn_dense.inp_dim\n",
    "                out_dim = nn_dense.out_dim\n",
    "\n",
    "                self.net = nn.Sequential(nn.Linear(state_dim, inp_dim), nn.ReLU(),\n",
    "                                         nn_dense,\n",
    "                                         nn.Linear(out_dim, 1), )\n",
    "            else:\n",
    "                self.net = nn.Sequential(nn.Linear(state_dim, mid_dim), nn.ReLU(),\n",
    "                                         nn.Linear(mid_dim, mid_dim), nn.ReLU(),\n",
    "                                         nn.Linear(mid_dim, mid_dim), Hardswish(),\n",
    "                                         nn.Linear(mid_dim, 1), )\n",
    "\n",
    "            self.net = nn.Sequential(nn.Linear(state_dim, mid_dim), nn.ReLU(),\n",
    "                                     nn.Linear(mid_dim, mid_dim), nn.ReLU(),\n",
    "                                     nn.Linear(mid_dim, mid_dim), nn.ReLU(),  # nn.Hardswish(),\n",
    "                                     nn.Linear(mid_dim, 1))\n",
    "        else:\n",
    "            def set_dim(i):\n",
    "                return int(12 * 1.5 ** i)\n",
    "\n",
    "            self.net = nn.Sequential(NnReshape(*state_dim),  # -> [batch_size, 4, 96, 96]\n",
    "                                     nn.Conv2d(state_dim[0], set_dim(0), 4, 2, bias=True), nn.LeakyReLU(),\n",
    "                                     nn.Conv2d(set_dim(0), set_dim(1), 3, 2, bias=False), nn.ReLU(),\n",
    "                                     nn.Conv2d(set_dim(1), set_dim(2), 3, 2, bias=False), nn.ReLU(),\n",
    "                                     nn.Conv2d(set_dim(2), set_dim(3), 3, 2, bias=True), nn.ReLU(),\n",
    "                                     nn.Conv2d(set_dim(3), set_dim(4), 3, 1, bias=True), nn.ReLU(),\n",
    "                                     nn.Conv2d(set_dim(4), set_dim(5), 3, 1, bias=True), nn.ReLU(),\n",
    "                                     NnReshape(-1),\n",
    "                                     nn.Linear(set_dim(5), mid_dim), nn.ReLU(),\n",
    "                                     nn.Linear(mid_dim, 1))\n",
    "\n",
    "        layer_norm(self.net[-1], std=0.5)  # output layer for Q value\n",
    "\n",
    "    def forward(self, state):\n",
    "        return self.net(state)  # Q value\n",
    "\n",
    "\n",
    "class CriticTwin(nn.Module):\n",
    "    def __init__(self, mid_dim, state_dim, action_dim, if_use_dn=False):\n",
    "        super().__init__()\n",
    "\n",
    "        if if_use_dn:  # use DenseNet (DenseNet has both shallow and deep linear layer)\n",
    "            nn_dense = DenseNet(mid_dim // 2)\n",
    "            inp_dim = nn_dense.inp_dim\n",
    "            out_dim = nn_dense.out_dim\n",
    "\n",
    "            self.net_sa = nn.Sequential(nn.Linear(state_dim + action_dim, inp_dim), nn.ReLU(),\n",
    "                                        nn_dense, )  # state-action value function\n",
    "        else:  # use a simple network for actor. Deeper network does not mean better performance in RL.\n",
    "            self.net_sa = nn.Sequential(nn.Linear(state_dim + action_dim, mid_dim), nn.ReLU(),\n",
    "                                        nn.Linear(mid_dim, mid_dim), nn.ReLU())\n",
    "            out_dim = mid_dim\n",
    "\n",
    "        self.net_q1 = nn.Linear(out_dim, 1)\n",
    "        self.net_q2 = nn.Linear(out_dim, 1)\n",
    "        layer_norm(self.net_q1, std=0.1)\n",
    "        layer_norm(self.net_q2, std=0.1)\n",
    "\n",
    "    def forward(self, state, action):\n",
    "        tmp = self.net_sa(torch.cat((state, action), dim=1))\n",
    "        return self.net_q1(tmp)  # one Q value\n",
    "\n",
    "    def get_q1_q2(self, state, action):\n",
    "        tmp = self.net_sa(torch.cat((state, action), dim=1))\n",
    "        return self.net_q1(tmp), self.net_q2(tmp)  # two Q values\n",
    "\n",
    "\n",
    "'''Parameter sharing Network'''\n",
    "\n",
    "\n",
    "class SharedDPG(nn.Module):  # DPG means deterministic policy gradient\n",
    "    def __init__(self, state_dim, action_dim, mid_dim):\n",
    "        super().__init__()\n",
    "        nn_dense = DenseNet(mid_dim // 2)\n",
    "        inp_dim = nn_dense.inp_dim\n",
    "        out_dim = nn_dense.out_dim\n",
    "\n",
    "        self.enc_s = nn.Sequential(nn.Linear(state_dim, mid_dim), nn.ReLU(),\n",
    "                                   nn.Linear(mid_dim, inp_dim))\n",
    "        self.enc_a = nn.Sequential(nn.Linear(action_dim, mid_dim), nn.ReLU(),\n",
    "                                   nn.Linear(mid_dim, inp_dim))\n",
    "\n",
    "        self.net = nn_dense\n",
    "\n",
    "        self.dec_a = nn.Sequential(nn.Linear(out_dim, mid_dim), Hardswish(),\n",
    "                                   nn.Linear(mid_dim, action_dim), nn.Tanh())\n",
    "        self.dec_q = nn.Sequential(nn.Linear(out_dim, mid_dim), Hardswish(),\n",
    "                                   nn.utils.spectral_norm(nn.Linear(mid_dim, 1)))\n",
    "\n",
    "    @staticmethod\n",
    "    def add_noise(a, noise_std):\n",
    "        a_temp = torch.normal(a, noise_std)\n",
    "        mask = torch.tensor((a_temp < -1.0) + (a_temp > 1.0), dtype=torch.float32).cuda()\n",
    "\n",
    "        noise_uniform = torch.rand_like(a)\n",
    "        a_noise = noise_uniform * mask + a_temp * (-mask + 1)\n",
    "        return a_noise\n",
    "\n",
    "    def forward(self, s, noise_std=0.0):  # actor\n",
    "        s_ = self.enc_s(s)\n",
    "        a_ = self.net(s_)\n",
    "        a = self.dec_a(a_)\n",
    "        return a if noise_std == 0.0 else self.add_noise(a, noise_std)\n",
    "\n",
    "    def critic(self, s, a):\n",
    "        s_ = self.enc_s(s)\n",
    "        a_ = self.enc_a(a)\n",
    "        q_ = self.net(s_ + a_)\n",
    "        q = self.dec_q(q_)\n",
    "        return q\n",
    "\n",
    "    def next_q_action(self, s, s_next, noise_std):\n",
    "        s_ = self.enc_s(s)\n",
    "        a_ = self.net(s_)\n",
    "        a = self.dec_a(a_)\n",
    "\n",
    "        '''q_target (without noise)'''\n",
    "        a_ = self.enc_a(a)\n",
    "        s_next_ = self.enc_s(s_next)\n",
    "        q_target0_ = self.net(s_next_ + a_)\n",
    "        q_target0 = self.dec_q(q_target0_)\n",
    "\n",
    "        '''q_target (with noise)'''\n",
    "        a_noise = self.add_noise(a, noise_std)\n",
    "        a_noise_ = self.enc_a(a_noise)\n",
    "        q_target1_ = self.net(s_next_ + a_noise_)\n",
    "        q_target1 = self.dec_q(q_target1_)\n",
    "\n",
    "        q_target = (q_target0 + q_target1) * 0.5\n",
    "        return q_target, a\n",
    "\n",
    "\n",
    "class SharedSPG(nn.Module):  # SPG means stochastic policy gradient\n",
    "    def __init__(self, mid_dim, state_dim, action_dim):\n",
    "        super().__init__()\n",
    "        self.log_sqrt_2pi_sum = np.log(np.sqrt(2 * np.pi)) * action_dim\n",
    "        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "        nn_dense = DenseNet(mid_dim // 2)\n",
    "        inp_dim = nn_dense.inp_dim\n",
    "        out_dim = nn_dense.out_dim\n",
    "\n",
    "        self.enc_s = nn.Sequential(nn.Linear(state_dim, mid_dim), nn.ReLU(),\n",
    "                                   nn.Linear(mid_dim, inp_dim), )  # state\n",
    "        self.enc_a = nn.Sequential(nn.Linear(action_dim, mid_dim), nn.ReLU(),\n",
    "                                   nn.Linear(mid_dim, inp_dim), )  # action without nn.Tanh()\n",
    "\n",
    "        self.net = nn_dense\n",
    "\n",
    "        self.dec_a = nn.Sequential(nn.Linear(out_dim, mid_dim), nn.ReLU(),\n",
    "                                   nn.Linear(mid_dim, action_dim), )  # action_mean\n",
    "        self.dec_d = nn.Sequential(nn.Linear(out_dim, mid_dim), nn.ReLU(),\n",
    "                                   nn.Linear(mid_dim, action_dim), )  # action_std_log (d means standard deviation)\n",
    "        self.dec_q1 = nn.Sequential(nn.Linear(out_dim, mid_dim), nn.ReLU(),\n",
    "                                    nn.Linear(mid_dim, 1), )  # q1 value\n",
    "        self.dec_q2 = nn.Sequential(nn.Linear(out_dim, mid_dim), nn.ReLU(),\n",
    "                                    nn.Linear(mid_dim, 1), )  # q2 value\n",
    "\n",
    "        layer_norm(self.dec_a[-1], std=0.5)\n",
    "        layer_norm(self.dec_d[-1], std=0.1)\n",
    "        layer_norm(self.dec_q1[-1], std=0.5)\n",
    "        layer_norm(self.dec_q2[-1], std=0.5)\n",
    "\n",
    "    def forward(self, s):\n",
    "        x = self.enc_s(s)\n",
    "        x = self.net(x)\n",
    "        a_avg = self.dec_a(x)\n",
    "        return a_avg.tanh()\n",
    "\n",
    "    def get_noise_action(self, s):\n",
    "        s_ = self.enc_s(s)\n",
    "        a_ = self.net(s_)\n",
    "        a_avg = self.dec_a(a_)  # NOTICE! it is a_avg without tensor.tanh()\n",
    "\n",
    "        a_std_log = self.dec_d(a_).clamp(-20, 2)\n",
    "        a_std = a_std_log.exp()\n",
    "\n",
    "        action = torch.normal(a_avg, a_std)  # NOTICE! it is action without .tanh()\n",
    "        return action.tanh()\n",
    "\n",
    "    def get_a_logprob(self, state):  # actor\n",
    "        s_ = self.enc_s(state)\n",
    "        a_ = self.net(s_)\n",
    "\n",
    "        \"\"\"add noise to action, stochastic policy\"\"\"\n",
    "        a_avg = self.dec_a(a_)  # NOTICE! it is action without .tanh()\n",
    "        a_std_log = self.dec_d(a_).clamp(-20, 2)\n",
    "        a_std = a_std_log.exp()\n",
    "\n",
    "        noise = torch.randn_like(a_avg, requires_grad=True)\n",
    "        a_noise = a_avg + a_std * noise\n",
    "\n",
    "        a_noise_tanh = a_noise.tanh()\n",
    "        fix_term = (-a_noise_tanh.pow(2) + 1.00001).log()\n",
    "        logprob = (noise.pow(2) / 2 + a_std_log + fix_term).sum(1, keepdim=True) + self.log_sqrt_2pi_sum\n",
    "        return a_noise_tanh, logprob\n",
    "\n",
    "    def get_q_logprob(self, state):\n",
    "        s_ = self.enc_s(state)\n",
    "        a_ = self.net(s_)\n",
    "\n",
    "        \"\"\"add noise to action, stochastic policy\"\"\"\n",
    "        a_avg = self.dec_a(a_)  # NOTICE! it is action without .tanh()\n",
    "        a_std_log = self.dec_d(a_).clamp(-20, 2)\n",
    "        a_std = a_std_log.exp()\n",
    "\n",
    "        noise = torch.randn_like(a_avg, requires_grad=True)\n",
    "        a_noise = a_avg + a_std * noise\n",
    "\n",
    "        a_noise_tanh = a_noise.tanh()\n",
    "        fix_term = (-a_noise_tanh.pow(2) + 1.00001).log()\n",
    "        logprob = (noise.pow(2) / 2 + a_std_log + fix_term).sum(1, keepdim=True) + self.log_sqrt_2pi_sum\n",
    "\n",
    "        '''get q'''\n",
    "        a_ = self.enc_a(a_noise_tanh)\n",
    "        q_ = self.net(s_ + a_)\n",
    "        q = torch.min(self.dec_q1(q_), self.dec_q2(q_))\n",
    "        return q, logprob\n",
    "\n",
    "    def get_q1_q2(self, s, a):  # critic\n",
    "        s_ = self.enc_s(s)\n",
    "        a_ = self.enc_a(a)\n",
    "        q_ = self.net(s_ + a_)\n",
    "        q1 = self.dec_q1(q_)\n",
    "        q2 = self.dec_q2(q_)\n",
    "        return q1, q2\n",
    "\n",
    "\n",
    "class SharedPPO(nn.Module):  # Pixel-level state version\n",
    "    def __init__(self, state_dim, action_dim, mid_dim):\n",
    "        super().__init__()\n",
    "\n",
    "        def set_dim(i):\n",
    "            return int(12 * 1.5 ** i)\n",
    "\n",
    "        if isinstance(state_dim, int):\n",
    "            self.enc_s = nn.Sequential(nn.Linear(state_dim, mid_dim), nn.ReLU(),\n",
    "                                       nn.Linear(mid_dim, mid_dim))  # the only difference.\n",
    "        else:\n",
    "            self.enc_s = nn.Sequential(NnReshape(*state_dim),  # -> [batch_size, 4, 96, 96]\n",
    "                                       nn.Conv2d(state_dim[0], set_dim(0), 4, 2, bias=True), nn.LeakyReLU(),\n",
    "                                       nn.Conv2d(set_dim(0), set_dim(1), 3, 2, bias=False), nn.ReLU(),\n",
    "                                       nn.Conv2d(set_dim(1), set_dim(2), 3, 2, bias=False), nn.ReLU(),\n",
    "                                       nn.Conv2d(set_dim(2), set_dim(3), 3, 2, bias=True), nn.ReLU(),\n",
    "                                       nn.Conv2d(set_dim(3), set_dim(4), 3, 1, bias=True), nn.ReLU(),\n",
    "                                       nn.Conv2d(set_dim(4), set_dim(5), 3, 1, bias=True), nn.ReLU(),\n",
    "                                       NnReshape(-1),\n",
    "                                       nn.Linear(set_dim(5), mid_dim), nn.ReLU(),\n",
    "                                       nn.Linear(mid_dim, mid_dim))\n",
    "        out_dim = mid_dim\n",
    "\n",
    "        self.dec_a = nn.Sequential(nn.Linear(out_dim, mid_dim), nn.ReLU(),\n",
    "                                   nn.Linear(mid_dim, action_dim))\n",
    "        self.a_std_log = nn.Parameter(torch.zeros(1, action_dim) - 0.5, requires_grad=True)\n",
    "\n",
    "        self.dec_q1 = nn.Sequential(nn.Linear(out_dim, mid_dim), nn.ReLU(),\n",
    "                                    nn.Linear(mid_dim, 1))\n",
    "        self.dec_q2 = nn.Sequential(nn.Linear(out_dim, mid_dim), nn.ReLU(),\n",
    "                                    nn.Linear(mid_dim, 1))\n",
    "\n",
    "        layer_norm(self.dec_a[-1], std=0.01)\n",
    "        layer_norm(self.dec_q1[-1], std=0.01)\n",
    "        layer_norm(self.dec_q2[-1], std=0.01)\n",
    "\n",
    "        self.sqrt_2pi_log = np.log(np.sqrt(2 * np.pi))\n",
    "\n",
    "    def forward(self, s):\n",
    "        s_ = self.enc_s(s)\n",
    "        a_avg = self.dec_a(s_)\n",
    "        return a_avg.tanh()\n",
    "\n",
    "    def get_action_noise(self, state):\n",
    "        s_ = self.enc_s(state)\n",
    "        a_avg = self.dec_a(s_)\n",
    "        a_std = self.a_std_log.exp()\n",
    "\n",
    "        # a_noise = torch.normal(a_avg, a_std) # same as below\n",
    "        noise = torch.randn_like(a_avg)\n",
    "        a_noise = a_avg + noise * a_std\n",
    "        return a_noise, noise\n",
    "\n",
    "    def get_q_logprob(self, state, noise):\n",
    "        s_ = self.enc_s(state)\n",
    "\n",
    "        q = torch.min(self.dec_q1(s_), self.dec_q2(s_))\n",
    "        logprob = -(noise.pow(2) / 2 + self.a_std_log + self.sqrt_2pi_log).sum(1)\n",
    "        return q, logprob\n",
    "\n",
    "    def get_q1_q2_logprob(self, state, action):\n",
    "        s_ = self.enc_s(state)\n",
    "\n",
    "        q1 = self.dec_q1(s_)\n",
    "        q2 = self.dec_q2(s_)\n",
    "\n",
    "        a_avg = self.dec_a(s_)\n",
    "        a_std = self.a_std_log.exp()\n",
    "        logprob = -(((a_avg - action) / a_std).pow(2) / 2 + self.a_std_log + self.sqrt_2pi_log).sum(1)\n",
    "        return q1, q2, logprob\n",
    "\n",
    "\n",
    "\"\"\"utils\"\"\"\n",
    "\n",
    "\n",
    "class NnReshape(nn.Module):\n",
    "    def __init__(self, *args):\n",
    "        super().__init__()\n",
    "        self.args = args\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x.view((x.size(0),) + self.args)\n",
    "\n",
    "\n",
    "class DenseNet(nn.Module):  # plan to hyper-param: layer_number\n",
    "    def __init__(self, lay_dim):\n",
    "        super().__init__()\n",
    "        self.dense1 = nn.Sequential(nn.Linear(lay_dim * 1, lay_dim * 1), Hardswish())\n",
    "        self.dense2 = nn.Sequential(nn.Linear(lay_dim * 2, lay_dim * 2), Hardswish())\n",
    "        self.inp_dim = lay_dim\n",
    "        self.out_dim = lay_dim * 4\n",
    "\n",
    "    def forward(self, x1):  # x1.shape==(-1, lay_dim*1)\n",
    "        x2 = torch.cat((x1, self.dense1(x1)), dim=1)\n",
    "        x3 = torch.cat((x2, self.dense2(x2)), dim=1)\n",
    "        return x3  # x2.shape==(-1, lay_dim*4)\n",
    "\n",
    "\n",
    "class ConcatNet(nn.Module):  # concatenate\n",
    "    def __init__(self, lay_dim):\n",
    "        super().__init__()\n",
    "        self.dense1 = nn.Sequential(nn.Linear(lay_dim, lay_dim), nn.ReLU(),\n",
    "                                    nn.Linear(lay_dim, lay_dim), Hardswish(), )\n",
    "        self.dense2 = nn.Sequential(nn.Linear(lay_dim, lay_dim), nn.ReLU(),\n",
    "                                    nn.Linear(lay_dim, lay_dim), Hardswish(), )\n",
    "        self.dense3 = nn.Sequential(nn.Linear(lay_dim, lay_dim), nn.ReLU(),\n",
    "                                    nn.Linear(lay_dim, lay_dim), Hardswish(), )\n",
    "        self.dense4 = nn.Sequential(nn.Linear(lay_dim, lay_dim), nn.ReLU(),\n",
    "                                    nn.Linear(lay_dim, lay_dim), Hardswish(), )\n",
    "        self.inp_dim = lay_dim\n",
    "        self.out_dim = lay_dim * 4\n",
    "\n",
    "    def forward(self, x0):\n",
    "        x1 = self.dense1(x0)\n",
    "        x2 = self.dense2(x0)\n",
    "        x3 = self.dense3(x0)\n",
    "        x4 = self.dense4(x0)\n",
    "\n",
    "        return torch.cat((x1, x2, x3, x4), dim=1)\n",
    "\n",
    "\n",
    "def layer_norm(layer, std=1.0, bias_const=1e-6):\n",
    "    torch.nn.init.orthogonal_(layer.weight, std)\n",
    "    torch.nn.init.constant_(layer.bias, bias_const)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Utils'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "import torch\n",
    "import numpy as np\n",
    "import numpy.random as rd\n",
    "from copy import deepcopy\n",
    "\n",
    "\"\"\"[ElegantRL](https://github.com/AI4Finance-LLC/ElegantRL)\"\"\"\n",
    "\n",
    "\n",
    "class AgentBase:\n",
    "    def __init__(self):\n",
    "        self.learning_rate = 1e-4\n",
    "        self.soft_update_tau = 2 ** -8  # 5e-3 ~= 2 ** -8\n",
    "        self.state = None  # set for self.update_buffer(), initialize before training\n",
    "        self.device = None\n",
    "\n",
    "        self.act = self.act_target = None\n",
    "        self.cri = self.cri_target = None\n",
    "        self.act_optimizer = None\n",
    "        self.cri_optimizer = None\n",
    "        self.criterion = None\n",
    "        self.get_obj_critic = None\n",
    "\n",
    "    def init(self, net_dim, state_dim, action_dim, if_per=False):\n",
    "        \"\"\"initialize the self.object in `__init__()`\n",
    "        replace by different DRL algorithms\n",
    "        explict call self.init() for multiprocessing.\n",
    "        `int net_dim` the dimension of networks (the width of neural networks)\n",
    "        `int state_dim` the dimension of state (the number of state vector)\n",
    "        `int action_dim` the dimension of action (the number of discrete action)\n",
    "        `bool if_per` Prioritized Experience Replay for sparse reward\n",
    "        \"\"\"\n",
    "\n",
    "    def select_action(self, state) -> np.ndarray:\n",
    "        \"\"\"Select actions for exploration\n",
    "        :array state: state.shape==(state_dim, )\n",
    "        :return array action: action.shape==(action_dim, ), (action.min(), action.max())==(-1, +1)\n",
    "        \"\"\"\n",
    "        states = torch.as_tensor((state,), dtype=torch.float32, device=self.device).detach_()\n",
    "        action = self.act(states)[0]\n",
    "        return action.cpu().numpy()\n",
    "\n",
    "    def explore_env(self, env, buffer, target_step, reward_scale, gamma) -> int:\n",
    "        \"\"\"actor explores in env, then stores the env transition to ReplayBuffer\n",
    "        :env: RL training environment. env.reset() env.step()\n",
    "        :buffer: Experience Replay Buffer.\n",
    "        :int target_step: explored target_step number of step in env\n",
    "        :float reward_scale: scale reward, 'reward * reward_scale'\n",
    "        :float gamma: discount factor, 'mask = 0.0 if done else gamma'\n",
    "        :return int target_step: collected target_step number of step in env\n",
    "        \"\"\"\n",
    "        for _ in range(target_step):\n",
    "            action = self.select_action(self.state)\n",
    "            next_s, reward, done, _ = env.step(action)\n",
    "            other = (reward * reward_scale, 0.0 if done else gamma, *action)\n",
    "            buffer.append_buffer(self.state, other)\n",
    "            self.state = env.reset() if done else next_s\n",
    "        return target_step\n",
    "\n",
    "    def update_net(self, buffer, target_step, batch_size, repeat_times) -> (float, float):\n",
    "        \"\"\"update the neural network by sampling batch data from ReplayBuffer\n",
    "        replace by different DRL algorithms.\n",
    "        return the objective value as training information to help fine-tuning\n",
    "        `buffer` Experience replay buffer.\n",
    "        :int target_step: explore target_step number of step in env\n",
    "        `int batch_size` sample batch_size of data for Stochastic Gradient Descent\n",
    "        :float repeat_times: the times of sample batch = int(target_step * repeat_times) in off-policy\n",
    "        :return float obj_a: the objective value of actor\n",
    "        :return float obj_c: the objective value of critic\n",
    "        \"\"\"\n",
    "\n",
    "    def save_load_model(self, cwd, if_save):\n",
    "        \"\"\"save or load model files\n",
    "        :str cwd: current working directory, we save model file here\n",
    "        :bool if_save: save model or load model\n",
    "        \"\"\"\n",
    "        act_save_path = '{}/actor.pth'.format(cwd)\n",
    "        cri_save_path = '{}/critic.pth'.format(cwd)\n",
    "\n",
    "        def load_torch_file(network, save_path):\n",
    "            network_dict = torch.load(save_path, map_location=lambda storage, loc: storage)\n",
    "            network.load_state_dict(network_dict)\n",
    "\n",
    "        if if_save:\n",
    "            if self.act is not None:\n",
    "                torch.save(self.act.state_dict(), act_save_path)\n",
    "            if self.cri is not None:\n",
    "                torch.save(self.cri.state_dict(), cri_save_path)\n",
    "        elif (self.act is not None) and os.path.exists(act_save_path):\n",
    "            load_torch_file(self.act, act_save_path)\n",
    "            print(\"Loaded act:\", cwd)\n",
    "        elif (self.cri is not None) and os.path.exists(cri_save_path):\n",
    "            load_torch_file(self.cri, cri_save_path)\n",
    "            print(\"Loaded cri:\", cwd)\n",
    "        else:\n",
    "            print(\"FileNotFound when load_model: {}\".format(cwd))\n",
    "\n",
    "    @staticmethod\n",
    "    def soft_update(target_net, current_net, tau):\n",
    "        \"\"\"soft update a target network via current network\n",
    "        :nn.Module target_net: target network update via a current network, it is more stable\n",
    "        :nn.Module current_net: current network update via an optimizer\n",
    "        \"\"\"\n",
    "        for tar, cur in zip(target_net.parameters(), current_net.parameters()):\n",
    "            tar.data.copy_(cur.data.__mul__(tau) + tar.data.__mul__(1 - tau))\n",
    "\n",
    "\n",
    "'''Value-based Methods (DQN variances)'''\n",
    "\n",
    "\n",
    "class AgentDQN(AgentBase):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.explore_rate = 0.1  # the probability of choosing action randomly in epsilon-greedy\n",
    "        self.action_dim = None  # chose discrete action randomly in epsilon-greedy\n",
    "\n",
    "    def init(self, net_dim, state_dim, action_dim, if_per=False):\n",
    "        self.action_dim = action_dim\n",
    "        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "        self.cri = QNet(net_dim, state_dim, action_dim).to(self.device)\n",
    "        self.cri_target = deepcopy(self.cri)\n",
    "        self.cri_optimizer = torch.optim.Adam(self.cri.parameters(), lr=self.learning_rate)\n",
    "        self.act = self.cri  # to keep the same from Actor-Critic framework\n",
    "\n",
    "        self.criterion = torch.nn.MSELoss(reduction='none' if if_per else 'mean')\n",
    "        if if_per:\n",
    "            self.get_obj_critic = self.get_obj_critic_per\n",
    "        else:\n",
    "            self.get_obj_critic = self.get_obj_critic_raw\n",
    "\n",
    "    def select_action(self, state) -> int:  # for discrete action space\n",
    "        if rd.rand() < self.explore_rate:  # epsilon-greedy\n",
    "            a_int = rd.randint(self.action_dim)  # choosing action randomly\n",
    "        else:\n",
    "            states = torch.as_tensor((state,), dtype=torch.float32, device=self.device).detach_()\n",
    "            action = self.act(states)[0]\n",
    "            a_int = action.argmax(dim=0).cpu().numpy()\n",
    "        return a_int\n",
    "\n",
    "    def explore_env(self, env, buffer, target_step, reward_scale, gamma) -> int:\n",
    "        for _ in range(target_step):\n",
    "            action = self.select_action(self.state)\n",
    "            next_s, reward, done, _ = env.step(action)\n",
    "\n",
    "            other = (reward * reward_scale, 0.0 if done else gamma, action)  # action is an int\n",
    "            buffer.append_buffer(self.state, other)\n",
    "            self.state = env.reset() if done else next_s\n",
    "        return target_step\n",
    "\n",
    "    def update_net(self, buffer, target_step, batch_size, repeat_times) -> (float, float):\n",
    "        buffer.update_now_len_before_sample()\n",
    "\n",
    "        q_value = obj_critic = None\n",
    "        for _ in range(int(target_step * repeat_times)):\n",
    "            obj_critic, q_value = self.get_obj_critic(buffer, batch_size)\n",
    "\n",
    "            self.cri_optimizer.zero_grad()\n",
    "            obj_critic.backward()\n",
    "            self.cri_optimizer.step()\n",
    "            self.soft_update(self.cri_target, self.cri, self.soft_update_tau)\n",
    "        return q_value.mean().item(), obj_critic.item()\n",
    "\n",
    "    def get_obj_critic_raw(self, buffer, batch_size):\n",
    "        with torch.no_grad():\n",
    "            reward, mask, action, state, next_s = buffer.sample_batch(batch_size)\n",
    "            next_q = self.cri_target(next_s).max(dim=1, keepdim=True)[0]\n",
    "            q_label = reward + mask * next_q\n",
    "\n",
    "        q_value = self.cri(state).gather(1, action.type(torch.long))\n",
    "        obj_critic = self.criterion(q_value, q_label)\n",
    "        return obj_critic, q_value\n",
    "\n",
    "    def get_obj_critic_per(self, buffer, batch_size):\n",
    "        with torch.no_grad():\n",
    "            reward, mask, action, state, next_s, is_weights = buffer.sample_batch(batch_size)\n",
    "            next_q = self.cri_target(next_s).max(dim=1, keepdim=True)[0]\n",
    "            q_label = reward + mask * next_q\n",
    "\n",
    "        q_value = self.cri(state).gather(1, action.type(torch.long))\n",
    "        obj_critic = (self.criterion(q_value, q_label) * is_weights).mean()\n",
    "        return obj_critic, q_value\n",
    "\n",
    "\n",
    "class AgentDuelingDQN(AgentDQN):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.explore_rate = 0.25  # the probability of choosing action randomly in epsilon-greedy\n",
    "\n",
    "    def init(self, net_dim, state_dim, action_dim, if_per=False):\n",
    "        self.action_dim = action_dim\n",
    "        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "        self.cri = QNetDuel(net_dim, state_dim, action_dim).to(self.device)\n",
    "        self.cri_target = deepcopy(self.cri)\n",
    "        self.act = self.cri\n",
    "\n",
    "        self.cri_optimizer = torch.optim.Adam(self.cri.parameters(), lr=self.learning_rate)\n",
    "        self.criterion = torch.nn.MSELoss(reduction='none' if if_per else 'mean')\n",
    "        self.get_obj_critic = self.get_obj_critic_per if if_per else self.get_obj_critic_raw\n",
    "\n",
    "\n",
    "class AgentDoubleDQN(AgentDQN):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.explore_rate = 0.25  # the probability of choosing action randomly in epsilon-greedy\n",
    "        self.softmax = torch.nn.Softmax(dim=1)\n",
    "\n",
    "    def init(self, net_dim, state_dim, action_dim, if_per=False):\n",
    "        self.action_dim = action_dim\n",
    "        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "        self.cri = QNetTwin(net_dim, state_dim, action_dim).to(self.device)\n",
    "        self.cri_target = deepcopy(self.cri)\n",
    "        self.cri_optimizer = torch.optim.Adam(self.cri.parameters(), lr=self.learning_rate)\n",
    "\n",
    "        self.act = self.cri\n",
    "        self.criterion = torch.nn.SmoothL1Loss(reduction='none' if if_per else 'mean')\n",
    "        self.get_obj_critic = self.get_obj_critic_per if if_per else self.get_obj_critic_raw\n",
    "\n",
    "    def select_action(self, state) -> int:  # for discrete action space\n",
    "        states = torch.as_tensor((state,), dtype=torch.float32, device=self.device).detach_()\n",
    "        actions = self.act(states)\n",
    "        if rd.rand() < self.explore_rate:  # epsilon-greedy\n",
    "            action = self.softmax(actions)[0]\n",
    "            a_prob = action.detach().cpu().numpy()  # choose action according to Q value\n",
    "            a_int = rd.choice(self.action_dim, p=a_prob)\n",
    "        else:\n",
    "            action = actions[0]\n",
    "            a_int = action.argmax(dim=0).cpu().numpy()\n",
    "        return a_int\n",
    "\n",
    "    def get_obj_critic_raw(self, buffer, batch_size):\n",
    "        with torch.no_grad():\n",
    "            reward, mask, action, state, next_s = buffer.sample_batch(batch_size)\n",
    "            next_q = torch.min(*self.cri_target.get_q1_q2(next_s))\n",
    "            next_q = next_q.max(dim=1, keepdim=True)[0]\n",
    "            q_label = reward + mask * next_q\n",
    "        act_int = action.type(torch.long)\n",
    "        q1, q2 = [qs.gather(1, act_int) for qs in self.act.get_q1_q2(state)]\n",
    "        obj_critic = self.criterion(q1, q_label) + self.criterion(q2, q_label)\n",
    "        return obj_critic, q1\n",
    "\n",
    "    def get_obj_critic_per(self, buffer, batch_size):\n",
    "        with torch.no_grad():\n",
    "            reward, mask, action, state, next_s, is_weights = buffer.sample_batch(batch_size)\n",
    "            next_q = torch.min(*self.cri_target.get_q1_q2(next_s))\n",
    "            next_q = next_q.max(dim=1, keepdim=True)[0]\n",
    "            q_label = reward + mask * next_q\n",
    "        act_int = action.type(torch.long)\n",
    "        q1, q2 = [qs.gather(1, act_int) for qs in self.act.get_q1_q2(state)]\n",
    "        obj_critic = ((self.criterion(q1, q_label) + self.criterion(q2, q_label)) * is_weights).mean()\n",
    "        return obj_critic, q1\n",
    "\n",
    "\n",
    "class AgentD3QN(AgentDoubleDQN):  # D3QN: Dueling Double DQN\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "    def init(self, net_dim, state_dim, action_dim, if_per=False):\n",
    "        self.action_dim = action_dim\n",
    "        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "        self.cri = QNetTwinDuel(net_dim, state_dim, action_dim).to(self.device)\n",
    "        self.cri_target = deepcopy(self.cri)\n",
    "        self.act = self.cri\n",
    "\n",
    "        self.criterion = torch.nn.SmoothL1Loss(reduction='none') if if_per else torch.nn.SmoothL1Loss()\n",
    "        self.cri_optimizer = torch.optim.Adam(self.cri.parameters(), lr=self.learning_rate)\n",
    "        self.criterion = torch.nn.SmoothL1Loss(reduction='none' if if_per else 'mean')\n",
    "        self.get_obj_critic = self.get_obj_critic_per if if_per else self.get_obj_critic_raw\n",
    "\n",
    "\n",
    "'''Actor-Critic Methods (Policy Gradient)'''\n",
    "\n",
    "\n",
    "class AgentDDPG(AgentBase):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.ou_explore_noise = 0.3  # explore noise of action\n",
    "        self.ou_noise = None\n",
    "\n",
    "    def init(self, net_dim, state_dim, action_dim, if_per=False):\n",
    "        self.ou_noise = OrnsteinUhlenbeckNoise(size=action_dim, sigma=self.ou_explore_noise)\n",
    "        # I don't recommend to use OU-Noise\n",
    "        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "        self.cri = Critic(net_dim, state_dim, action_dim).to(self.device)\n",
    "        self.cri_target = deepcopy(self.cri)\n",
    "        self.cri_optimizer = torch.optim.Adam(self.cri.parameters(), lr=self.learning_rate)\n",
    "\n",
    "        self.act = Actor(net_dim, state_dim, action_dim).to(self.device)\n",
    "        self.act_target = deepcopy(self.act)\n",
    "        self.act_optimizer = torch.optim.Adam(self.act.parameters(), lr=self.learning_rate)\n",
    "\n",
    "        self.criterion = torch.nn.SmoothL1Loss(reduction='none' if if_per else 'mean')\n",
    "        if if_per:\n",
    "            self.get_obj_critic = self.get_obj_critic_per\n",
    "        else:\n",
    "            self.get_obj_critic = self.get_obj_critic_raw\n",
    "\n",
    "    def select_action(self, state) -> np.ndarray:\n",
    "        states = torch.as_tensor((state,), dtype=torch.float32, device=self.device).detach_()\n",
    "        action = self.act(states)[0].cpu().numpy()\n",
    "        return (action + self.ou_noise()).clip(-1, 1)\n",
    "\n",
    "    def update_net(self, buffer, target_step, batch_size, repeat_times) -> (float, float):\n",
    "        buffer.update_now_len_before_sample()\n",
    "\n",
    "        obj_critic = obj_actor = None  # just for print return\n",
    "        for _ in range(int(target_step * repeat_times)):\n",
    "            obj_critic, state = self.get_obj_critic(buffer, batch_size)\n",
    "            self.cri_optimizer.zero_grad()\n",
    "            obj_critic.backward()\n",
    "            self.cri_optimizer.step()\n",
    "            self.soft_update(self.cri_target, self.cri, self.soft_update_tau)\n",
    "\n",
    "            q_value_pg = self.act(state)  # policy gradient\n",
    "            obj_actor = -self.cri_target(state, q_value_pg).mean()  # obj_actor\n",
    "            self.act_optimizer.zero_grad()\n",
    "            obj_actor.backward()\n",
    "            self.act_optimizer.step()\n",
    "            self.soft_update(self.act_target, self.act, self.soft_update_tau)\n",
    "        return obj_actor.item(), obj_critic.item()\n",
    "\n",
    "    def get_obj_critic_raw(self, buffer, batch_size):\n",
    "        with torch.no_grad():\n",
    "            reward, mask, action, state, next_s = buffer.sample_batch(batch_size)\n",
    "            next_q = self.cri_target(next_s, self.act_target(next_s))\n",
    "            q_label = reward + mask * next_q\n",
    "        q_value = self.cri(state, action)\n",
    "        obj_critic = self.criterion(q_value, q_label)\n",
    "        return obj_critic, state\n",
    "\n",
    "    def get_obj_critic_per(self, buffer, batch_size):\n",
    "        with torch.no_grad():\n",
    "            reward, mask, action, state, next_s, is_weights = buffer.sample_batch(batch_size)\n",
    "            next_q = self.cri_target(next_s, self.act_target(next_s))\n",
    "            q_label = reward + mask * next_q\n",
    "        q_value = self.cri(state, action)\n",
    "        obj_critic = (self.criterion(q_value, q_label) * is_weights).mean()\n",
    "\n",
    "        td_error = (q_label - q_value.detach()).abs()\n",
    "        buffer.td_error_update(td_error)\n",
    "        return obj_critic, state\n",
    "\n",
    "\n",
    "class AgentTD3(AgentDDPG):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.explore_noise = 0.1  # standard deviation of explore noise\n",
    "        self.policy_noise = 0.2  # standard deviation of policy noise\n",
    "        self.update_freq = 2  # delay update frequency, for soft target update\n",
    "\n",
    "    def init(self, net_dim, state_dim, action_dim, if_per=False):\n",
    "        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "        self.cri = CriticTwin(net_dim, state_dim, action_dim).to(self.device)\n",
    "        self.cri_target = deepcopy(self.cri)\n",
    "        self.cri_optimizer = torch.optim.Adam(self.cri.parameters(), lr=self.learning_rate)\n",
    "\n",
    "        self.act = Actor(net_dim, state_dim, action_dim).to(self.device)\n",
    "        self.act_target = deepcopy(self.act)\n",
    "        self.act_optimizer = torch.optim.Adam(self.act.parameters(), lr=self.learning_rate)\n",
    "\n",
    "        self.criterion = torch.nn.SmoothL1Loss(reduction='none' if if_per else 'mean')\n",
    "        if if_per:\n",
    "            self.get_obj_critic = self.get_obj_critic_per\n",
    "        else:\n",
    "            self.get_obj_critic = self.get_obj_critic_raw\n",
    "\n",
    "    def select_action(self, state) -> np.ndarray:\n",
    "        states = torch.as_tensor((state,), dtype=torch.float32, device=self.device).detach_()\n",
    "        action = self.act(states)[0]\n",
    "        action = (action + torch.randn_like(action) * self.explore_noise).clamp(-1, 1)\n",
    "        return action.cpu().numpy()\n",
    "\n",
    "    def update_net(self, buffer, target_step, batch_size, repeat_times) -> (float, float):\n",
    "        buffer.update_now_len_before_sample()\n",
    "\n",
    "        obj_critic = obj_actor = None\n",
    "        for i in range(int(target_step * repeat_times)):\n",
    "            obj_critic, state = self.get_obj_critic(buffer, batch_size)\n",
    "            self.cri_optimizer.zero_grad()\n",
    "            obj_critic.backward()\n",
    "            self.cri_optimizer.step()\n",
    "            if i % self.update_freq == 0:  # delay update\n",
    "                self.soft_update(self.cri_target, self.cri, self.soft_update_tau)\n",
    "\n",
    "            q_value_pg = self.act(state)  # policy gradient\n",
    "            obj_actor = -self.cri_target(state, q_value_pg).mean()  # obj_actor\n",
    "            self.act_optimizer.zero_grad()\n",
    "            obj_actor.backward()\n",
    "            self.act_optimizer.step()\n",
    "            if i % self.update_freq == 0:  # delay update\n",
    "                self.soft_update(self.act_target, self.act, self.soft_update_tau)\n",
    "\n",
    "        return obj_actor.item(), obj_critic.item() / 2\n",
    "\n",
    "    def get_obj_critic_raw(self, buffer, batch_size):\n",
    "        with torch.no_grad():\n",
    "            reward, mask, action, state, next_s = buffer.sample_batch(batch_size)\n",
    "            next_a = self.act_target.get_action(next_s, self.policy_noise)  # policy noise\n",
    "            next_q = torch.min(*self.cri_target.get_q1_q2(next_s, next_a))  # twin critics\n",
    "            q_label = reward + mask * next_q\n",
    "        q1, q2 = self.cri.get_q1_q2(state, action)\n",
    "        obj_critic = self.criterion(q1, q_label) + self.criterion(q2, q_label)  # twin critics\n",
    "        return obj_critic, state\n",
    "\n",
    "    def get_obj_critic_per(self, buffer, batch_size):\n",
    "        \"\"\"Prioritized Experience Replay\n",
    "        Contributor: Github GyChou\n",
    "        \"\"\"\n",
    "        with torch.no_grad():\n",
    "            reward, mask, action, state, next_s, is_weights = buffer.sample_batch(batch_size)\n",
    "            next_a = self.act_target.get_action(next_s, self.policy_noise)  # policy noise\n",
    "            next_q = torch.min(*self.cri_target.get_q1_q2(next_s, next_a))  # twin critics\n",
    "            q_label = reward + mask * next_q\n",
    "\n",
    "        q1, q2 = self.cri.get_q1_q2(state, action)\n",
    "        obj_critic = ((self.criterion(q1, q_label) + self.criterion(q2, q_label)) * is_weights).mean()\n",
    "\n",
    "        td_error = (q_label - torch.min(q1, q2).detach()).abs()\n",
    "        buffer.td_error_update(td_error)\n",
    "        return obj_critic, state\n",
    "\n",
    "\n",
    "class AgentSharedAC(AgentBase):  # use InterSAC instead of InterAC .Warning: sth. wrong with this code, need to check\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.explore_noise = 0.2  # standard deviation of explore noise\n",
    "        self.policy_noise = 0.4  # standard deviation of policy noise\n",
    "        self.update_freq = 2 ** 7  # delay update frequency, for hard target update\n",
    "        self.avg_loss_c = (-np.log(0.5)) ** 0.5  # old version reliable_lambda\n",
    "        self.optimizer = None\n",
    "\n",
    "    def init(self, net_dim, state_dim, action_dim, if_per=False):\n",
    "        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "        self.act = SharedDPG(state_dim, action_dim, net_dim).to(self.device)\n",
    "        self.act_target = deepcopy(self.act)\n",
    "\n",
    "        self.criterion = torch.nn.MSELoss(reduction='none') if if_per else torch.nn.MSELoss()\n",
    "        self.optimizer = torch.optim.Adam(self.act.parameters(), lr=self.learning_rate)\n",
    "\n",
    "    def update_net(self, buffer, target_step, batch_size, repeat_times) -> (float, float):\n",
    "        buffer.update_now_len_before_sample()\n",
    "\n",
    "        actor_obj = None  # just for print return\n",
    "\n",
    "        k = 1.0 + buffer.now_len / buffer.max_len\n",
    "        batch_size_ = int(batch_size * k)\n",
    "        update_times = int(target_step * k)\n",
    "\n",
    "        for i in range(update_times * repeat_times):\n",
    "            with torch.no_grad():\n",
    "                reward, mask, action, state, next_state = buffer.sample_batch(batch_size_)\n",
    "\n",
    "                next_q_label, next_action = self.act_target.next_q_action(state, next_state, self.policy_noise)\n",
    "                q_label = reward + mask * next_q_label\n",
    "\n",
    "            \"\"\"critic_obj\"\"\"\n",
    "            q_eval = self.act.critic(state, action)\n",
    "            critic_obj = self.criterion(q_eval, q_label)\n",
    "\n",
    "            '''auto reliable lambda'''\n",
    "            self.avg_loss_c = 0.995 * self.avg_loss_c + 0.005 * critic_obj.item() / 2  # soft update, twin critics\n",
    "            lamb = np.exp(-self.avg_loss_c ** 2)\n",
    "\n",
    "            '''actor correction term'''\n",
    "            actor_term = self.criterion(self.act(next_state), next_action)\n",
    "\n",
    "            if i % repeat_times == 0:\n",
    "                '''actor obj'''\n",
    "                action_pg = self.act(state)  # policy gradient\n",
    "                actor_obj = -self.act_target.critic(state, action_pg).mean()  # policy gradient\n",
    "                # NOTICE! It is very important to use act_target.critic here instead act.critic\n",
    "                # Or you can use act.critic.deepcopy(). Whatever you cannot use act.critic directly.\n",
    "\n",
    "                united_loss = critic_obj + actor_term * (1 - lamb) + actor_obj * (lamb * 0.5)\n",
    "            else:\n",
    "                united_loss = critic_obj + actor_term * (1 - lamb)\n",
    "\n",
    "            \"\"\"united loss\"\"\"\n",
    "            self.optimizer.zero_grad()\n",
    "            united_loss.backward()\n",
    "            self.optimizer.step()\n",
    "\n",
    "            if i % self.update_freq == self.update_freq and lamb > 0.1:\n",
    "                self.act_target.load_state_dict(self.act.state_dict())  # Hard Target Update\n",
    "\n",
    "        return actor_obj.item(), self.avg_loss_c\n",
    "\n",
    "\n",
    "class AgentSAC(AgentBase):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.target_entropy = None\n",
    "        self.alpha_log = None\n",
    "        self.alpha_optimizer = None\n",
    "        self.target_entropy = 1.0  # * np.log(action_dim)\n",
    "\n",
    "    def init(self, net_dim, state_dim, action_dim, if_per=False):\n",
    "        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "        self.target_entropy *= np.log(action_dim)\n",
    "        self.alpha_log = torch.tensor((-np.log(action_dim) * np.e,), dtype=torch.float32,\n",
    "                                      requires_grad=True, device=self.device)  # trainable parameter\n",
    "        self.alpha_optimizer = torch.optim.Adam((self.alpha_log,), self.learning_rate)\n",
    "\n",
    "        self.cri = CriticTwin(net_dim, state_dim, action_dim).to(self.device)\n",
    "        self.cri_target = deepcopy(self.cri)\n",
    "        self.cri_optimizer = torch.optim.Adam(self.cri.parameters(), lr=self.learning_rate)\n",
    "\n",
    "        self.act = ActorSAC(net_dim, state_dim, action_dim).to(self.device)\n",
    "        self.act_optimizer = torch.optim.Adam(self.act.parameters(), lr=self.learning_rate)\n",
    "\n",
    "        self.criterion = torch.nn.SmoothL1Loss(reduction='none' if if_per else 'mean')\n",
    "        if if_per:\n",
    "            self.get_obj_critic = self.get_obj_critic_per\n",
    "        else:\n",
    "            self.get_obj_critic = self.get_obj_critic_raw\n",
    "\n",
    "    def select_action(self, state) -> np.ndarray:\n",
    "        states = torch.as_tensor((state,), dtype=torch.float32, device=self.device).detach_()\n",
    "        action = self.act.get_action(states)[0]\n",
    "        return action.cpu().numpy()\n",
    "\n",
    "    def update_net(self, buffer, target_step, batch_size, repeat_times) -> (float, float):\n",
    "        buffer.update_now_len_before_sample()\n",
    "\n",
    "        alpha = self.alpha_log.exp().detach()\n",
    "        obj_critic = None\n",
    "        for _ in range(int(target_step * repeat_times)):\n",
    "            '''objective of critic'''\n",
    "            obj_critic, state = self.get_obj_critic(buffer, batch_size, alpha)\n",
    "            self.cri_optimizer.zero_grad()\n",
    "            obj_critic.backward()\n",
    "            self.cri_optimizer.step()\n",
    "            self.soft_update(self.cri_target, self.cri, self.soft_update_tau)\n",
    "\n",
    "            '''objective of alpha (temperature parameter automatic adjustment)'''\n",
    "            action_pg, logprob = self.act.get_action_logprob(state)  # policy gradient\n",
    "\n",
    "            obj_alpha = (self.alpha_log * (logprob - self.target_entropy).detach()).mean()\n",
    "            self.alpha_optimizer.zero_grad()\n",
    "            obj_alpha.backward()\n",
    "            self.alpha_optimizer.step()\n",
    "\n",
    "            '''objective of actor'''\n",
    "            alpha = self.alpha_log.exp().detach()\n",
    "            obj_actor = -(torch.min(*self.cri_target.get_q1_q2(state, action_pg)) + logprob * alpha).mean()\n",
    "\n",
    "            self.act_optimizer.zero_grad()\n",
    "            obj_actor.backward()\n",
    "            self.act_optimizer.step()\n",
    "\n",
    "        return obj_actor.item(), obj_critic.item()\n",
    "\n",
    "    def get_obj_critic_raw(self, buffer, batch_size, alpha):\n",
    "        with torch.no_grad():\n",
    "            reward, mask, action, state, next_s = buffer.sample_batch(batch_size)\n",
    "            next_a, next_logprob = self.act.get_action_logprob(next_s)\n",
    "            next_q = torch.min(*self.cri_target.get_q1_q2(next_s, next_a))\n",
    "            q_label = reward + mask * (next_q + next_logprob * alpha)\n",
    "        q1, q2 = self.cri.get_q1_q2(state, action)  # twin critics\n",
    "        obj_critic = self.criterion(q1, q_label) + self.criterion(q2, q_label)\n",
    "        return obj_critic, state\n",
    "\n",
    "    def get_obj_critic_per(self, buffer, batch_size, alpha):\n",
    "        with torch.no_grad():\n",
    "            reward, mask, action, state, next_s, is_weights = buffer.sample_batch(batch_size)\n",
    "            next_a, next_logprob = self.act.get_action_logprob(next_s)\n",
    "            next_q = torch.min(*self.cri_target.get_q1_q2(next_s, next_a))\n",
    "            q_label = reward + mask * (next_q + next_logprob * alpha)\n",
    "        q1, q2 = self.cri.get_q1_q2(state, action)  # twin critics\n",
    "        obj_critic = ((self.criterion(q1, q_label) + self.criterion(q2, q_label)) * is_weights).mean()\n",
    "\n",
    "        td_error = (q_label - torch.min(q1, q2).detach()).abs()\n",
    "        buffer.td_error_update(td_error)\n",
    "        return obj_critic, state\n",
    "\n",
    "\n",
    "class AgentModSAC(AgentSAC):  # Modified SAC using reliable_lambda and TTUR (Two Time-scale Update Rule)\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.if_use_dn = True\n",
    "        self.obj_c = (-np.log(0.5)) ** 0.5  # for reliable_lambda\n",
    "\n",
    "    def init(self, net_dim, state_dim, action_dim, if_per=False):\n",
    "        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "        self.target_entropy *= np.log(action_dim)\n",
    "        self.alpha_log = torch.tensor((-np.log(action_dim) * np.e,), dtype=torch.float32,\n",
    "                                      requires_grad=True, device=self.device)  # trainable parameter\n",
    "        self.alpha_optimizer = torch.optim.Adam((self.alpha_log,), self.learning_rate)\n",
    "\n",
    "        self.cri = CriticTwin(int(net_dim * 1.25), state_dim, action_dim, self.if_use_dn).to(self.device)\n",
    "        self.cri_target = deepcopy(self.cri)\n",
    "        self.cri_optimizer = torch.optim.Adam(self.cri.parameters(), self.learning_rate)\n",
    "\n",
    "        self.act = ActorSAC(net_dim, state_dim, action_dim, self.if_use_dn).to(self.device)\n",
    "        self.act_optimizer = torch.optim.Adam(self.act.parameters(), self.learning_rate)\n",
    "\n",
    "        self.criterion = torch.nn.SmoothL1Loss(reduction='none' if if_per else 'mean')\n",
    "        if if_per:\n",
    "            self.get_obj_critic = self.get_obj_critic_per\n",
    "        else:\n",
    "            self.get_obj_critic = self.get_obj_critic_raw\n",
    "\n",
    "    def update_net(self, buffer, target_step, batch_size, repeat_times) -> (float, float):\n",
    "        buffer.update_now_len_before_sample()\n",
    "\n",
    "        alpha = self.alpha_log.exp().detach()\n",
    "        update_a = 0\n",
    "        for update_c in range(1, int(buffer.now_len / batch_size * repeat_times)):\n",
    "            '''objective of critic (loss function of critic)'''\n",
    "            obj_critic, state = self.get_obj_critic(buffer, batch_size, alpha)\n",
    "            self.obj_c = 0.995 * self.obj_c + 0.0025 * obj_critic.item()  # for reliable_lambda\n",
    "            self.cri_optimizer.zero_grad()\n",
    "            obj_critic.backward()\n",
    "            self.cri_optimizer.step()\n",
    "            self.soft_update(self.cri_target, self.cri, self.soft_update_tau)\n",
    "\n",
    "            '''objective of actor using reliable_lambda and TTUR (Two Time-scales Update Rule)'''\n",
    "            reliable_lambda = np.exp(-self.obj_c ** 2)  # for reliable_lambda\n",
    "            if_update_a = (update_a / update_c) < (1 / (2 - reliable_lambda))\n",
    "            if if_update_a:  # auto TTUR\n",
    "                update_a += 1\n",
    "\n",
    "                '''objective of alpha (temperature parameter automatic adjustment)'''\n",
    "                action_pg, logprob = self.act.get_action_logprob(state)  # policy gradient\n",
    "\n",
    "                obj_alpha = (self.alpha_log * (logprob - self.target_entropy).detach()).mean() * reliable_lambda\n",
    "                self.alpha_optimizer.zero_grad()\n",
    "                obj_alpha.backward()\n",
    "                self.alpha_optimizer.step()\n",
    "\n",
    "                with torch.no_grad():\n",
    "                    self.alpha_log[:] = self.alpha_log.clamp(-20, 2)\n",
    "                alpha = self.alpha_log.exp().detach()\n",
    "\n",
    "                q_value_pg = torch.min(*self.cri_target.get_q1_q2(state, action_pg))\n",
    "                obj_actor = -(q_value_pg + logprob * alpha.detach()).mean() * reliable_lambda\n",
    "\n",
    "                self.act_optimizer.zero_grad()\n",
    "                obj_actor.backward()\n",
    "                self.act_optimizer.step()\n",
    "                self.soft_update(self.cri_target, self.cri, self.soft_update_tau)\n",
    "\n",
    "        return obj_actor.item(), self.obj_c\n",
    "\n",
    "\n",
    "class AgentSharedSAC(AgentSAC):  # Integrated Soft Actor-Critic\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.obj_c = (-np.log(0.5)) ** 0.5  # for reliable_lambda\n",
    "        self.optimizer = None\n",
    "\n",
    "    def init(self, net_dim, state_dim, action_dim, if_per=False):\n",
    "        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "        self.target_entropy *= np.log(action_dim)\n",
    "        self.alpha_log = torch.tensor((-np.log(action_dim) * np.e,), dtype=torch.float32,\n",
    "                                      requires_grad=True, device=self.device)  # trainable parameter\n",
    "\n",
    "        self.act = SharedSPG(net_dim, state_dim, action_dim).to(self.device)\n",
    "        self.act_target = deepcopy(self.act)\n",
    "\n",
    "        self.optimizer = torch.optim.Adam(\n",
    "            [{'params': self.act.enc_s.parameters(), 'lr': self.learning_rate * 0.9},  # more stable\n",
    "             {'params': self.act.enc_a.parameters(), },\n",
    "             {'params': self.act.net.parameters(), 'lr': self.learning_rate * 0.9},\n",
    "             {'params': self.act.dec_a.parameters(), },\n",
    "             {'params': self.act.dec_d.parameters(), },\n",
    "             {'params': self.act.dec_q1.parameters(), },\n",
    "             {'params': self.act.dec_q2.parameters(), },\n",
    "             {'params': (self.alpha_log,)}], lr=self.learning_rate)\n",
    "        self.criterion = torch.nn.SmoothL1Loss(reduction='none' if if_per else 'mean')\n",
    "        if if_per:\n",
    "            self.get_obj_critic = self.get_obj_critic_per\n",
    "        else:\n",
    "            self.get_obj_critic = self.get_obj_critic_raw\n",
    "\n",
    "    def select_action(self, state) -> np.ndarray:\n",
    "        states = torch.as_tensor((state,), dtype=torch.float32, device=self.device).detach_()\n",
    "        action = self.act.get_noise_action(states)[0]\n",
    "        return action.cpu().numpy()\n",
    "\n",
    "    def update_net(self, buffer, target_step, batch_size, repeat_times) -> (float, float):  # 1111\n",
    "        buffer.update_now_len_before_sample()\n",
    "\n",
    "        alpha = self.alpha_log.exp().detach()  # auto temperature parameter\n",
    "\n",
    "        k = 1.0 + buffer.now_len / buffer.max_len\n",
    "        batch_size_ = int(batch_size * k)  # increase batch_size\n",
    "        train_steps = int(target_step * k * repeat_times)  # increase training_step\n",
    "\n",
    "        update_a = 0\n",
    "        for update_c in range(1, train_steps):\n",
    "            '''objective of critic'''\n",
    "            obj_critic, state = self.get_obj_critic(buffer, batch_size_, alpha)\n",
    "            self.obj_c = 0.995 * self.obj_c + 0.005 * obj_critic.item() / 2  # soft update, twin critics\n",
    "            reliable_lambda = np.exp(-self.obj_c ** 2)\n",
    "\n",
    "            '''objective of alpha (temperature parameter automatic adjustment)'''\n",
    "            action_pg, logprob = self.act.get_a_logprob(state)  # policy gradient\n",
    "            obj_alpha = (self.alpha_log * (logprob - self.target_entropy).detach() * reliable_lambda).mean()\n",
    "\n",
    "            with torch.no_grad():\n",
    "                self.alpha_log[:] = self.alpha_log.clamp(-20, 2)\n",
    "                alpha = self.alpha_log.exp()  # .detach()\n",
    "\n",
    "            '''objective of actor using reliable_lambda and TTUR (Two Time-scales Update Rule)'''\n",
    "            if update_a / update_c < 1 / (2 - reliable_lambda):  # auto TTUR\n",
    "                update_a += 1\n",
    "                q_value_pg = torch.min(*self.act_target.get_q1_q2(state, action_pg)).mean()  # twin critics\n",
    "                obj_actor = -(q_value_pg + logprob * alpha.detach()).mean()  # policy gradient\n",
    "\n",
    "                obj_united = obj_critic + obj_alpha + obj_actor * reliable_lambda\n",
    "            else:\n",
    "                obj_united = obj_critic + obj_alpha\n",
    "\n",
    "            self.optimizer.zero_grad()\n",
    "            obj_united.backward()\n",
    "            self.optimizer.step()\n",
    "\n",
    "            self.soft_update(self.act_target, self.act, self.soft_update_tau)\n",
    "\n",
    "        return obj_actor.item(), self.obj_c\n",
    "\n",
    "\n",
    "class AgentPPO(AgentBase):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.ratio_clip = 0.3  # could be 0.2 ~ 0.5, ratio.clamp(1 - clip, 1 + clip),\n",
    "        self.lambda_entropy = 0.04  # could be 0.01 ~ 0.05\n",
    "        self.lambda_gae_adv = 0.97  # could be 0.95 ~ 0.99, GAE (Generalized Advantage Estimation. ICLR.2016.)\n",
    "        self.if_use_gae = False  # if use Generalized Advantage Estimation\n",
    "        self.if_on_policy = True  # AgentPPO is an on policy DRL algorithm\n",
    "        self.if_use_dn = False\n",
    "\n",
    "        self.noise = None\n",
    "        self.optimizer = None\n",
    "        self.compute_reward = None  # attribution\n",
    "\n",
    "    def init(self, net_dim, state_dim, action_dim, if_per=False):\n",
    "        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "        self.compute_reward = self.compute_reward_gae if self.if_use_gae else self.compute_reward_adv\n",
    "\n",
    "        self.cri = CriticAdv(state_dim, net_dim, self.if_use_dn).to(self.device)\n",
    "        self.act = ActorPPO(net_dim, state_dim, action_dim, self.if_use_dn).to(self.device)\n",
    "\n",
    "        self.optimizer = torch.optim.Adam([{'params': self.act.parameters(), 'lr': self.learning_rate},\n",
    "                                           {'params': self.cri.parameters(), 'lr': self.learning_rate}])\n",
    "        self.criterion = torch.nn.SmoothL1Loss()\n",
    "        assert if_per is False  # on-policy don't need PER\n",
    "\n",
    "    def select_action(self, state) -> tuple:\n",
    "        \"\"\"select action for PPO\n",
    "        :array state: state.shape==(state_dim, )\n",
    "        :return array action: state.shape==(action_dim, )\n",
    "        :return array noise: noise.shape==(action_dim, ), the noise\n",
    "        \"\"\"\n",
    "        states = torch.as_tensor((state,), dtype=torch.float32, device=self.device).detach()\n",
    "        actions, noises = self.act.get_action_noise(states)\n",
    "        return actions[0].detach().cpu().numpy(), noises[0].detach().cpu().numpy()  # todo remove detach()\n",
    "\n",
    "    def explore_env(self, env, buffer, target_step, reward_scale, gamma) -> int:\n",
    "        buffer.empty_buffer_before_explore()  # NOTICE! necessary for on-policy\n",
    "        # assert target_step == buffer.max_len - max_step\n",
    "\n",
    "        actual_step = 0\n",
    "        while actual_step < target_step:\n",
    "            state = env.reset()\n",
    "            for _ in range(env.max_step):\n",
    "                action, noise = self.select_action(state)\n",
    "\n",
    "                next_state, reward, done, _ = env.step(np.tanh(action))\n",
    "                actual_step += 1\n",
    "\n",
    "                other = (reward * reward_scale, 0.0 if done else gamma, *action, *noise)\n",
    "                buffer.append_buffer(state, other)\n",
    "                if done:\n",
    "                    break\n",
    "                state = next_state\n",
    "        return actual_step\n",
    "\n",
    "    def update_net(self, buffer, _target_step, batch_size, repeat_times=4) -> (float, float):\n",
    "        buffer.update_now_len_before_sample()\n",
    "        buf_len = buffer.now_len  # assert buf_len >= _target_step\n",
    "\n",
    "        '''Trajectory using reverse reward'''\n",
    "        with torch.no_grad():\n",
    "            buf_reward, buf_mask, buf_action, buf_noise, buf_state = buffer.sample_all()\n",
    "\n",
    "            bs = 2 ** 10  # set a smaller 'bs: batch size' when out of GPU memory.\n",
    "            buf_value = torch.cat([self.cri(buf_state[i:i + bs]) for i in range(0, buf_state.size(0), bs)], dim=0)\n",
    "            buf_logprob = -(buf_noise.pow(2).__mul__(0.5) + self.act.a_std_log + self.act.sqrt_2pi_log).sum(1)\n",
    "\n",
    "            buf_r_sum, buf_advantage = self.compute_reward(buf_len, buf_reward, buf_mask, buf_value)\n",
    "            del buf_reward, buf_mask, buf_noise\n",
    "\n",
    "        '''PPO: Surrogate objective of Trust Region'''\n",
    "        obj_critic = None\n",
    "        for _ in range(int(repeat_times * buf_len / batch_size)):\n",
    "            indices = torch.randint(buf_len, size=(batch_size,), requires_grad=False, device=self.device)\n",
    "\n",
    "            state = buf_state[indices]\n",
    "            action = buf_action[indices]\n",
    "            r_sum = buf_r_sum[indices]\n",
    "            logprob = buf_logprob[indices]\n",
    "            advantage = buf_advantage[indices]\n",
    "\n",
    "            new_logprob = self.act.compute_logprob(state, action)  # it is obj_actor\n",
    "            ratio = (new_logprob - logprob).exp()\n",
    "            obj_surrogate1 = advantage * ratio\n",
    "            obj_surrogate2 = advantage * ratio.clamp(1 - self.ratio_clip, 1 + self.ratio_clip)\n",
    "            obj_surrogate = -torch.min(obj_surrogate1, obj_surrogate2).mean()\n",
    "            obj_entropy = (new_logprob.exp() * new_logprob).mean()  # policy entropy\n",
    "            obj_actor = obj_surrogate + obj_entropy * self.lambda_entropy\n",
    "\n",
    "            value = self.cri(state).squeeze(1)  # critic network predicts the reward_sum (Q value) of state\n",
    "            obj_critic = self.criterion(value, r_sum)\n",
    "\n",
    "            obj_united = obj_actor + obj_critic / (r_sum.std() + 1e-5)\n",
    "            self.optimizer.zero_grad()\n",
    "            obj_united.backward()\n",
    "            self.optimizer.step()\n",
    "\n",
    "        return obj_actor.item(), obj_critic.item()\n",
    "\n",
    "    def compute_reward_adv(self, buf_len, buf_reward, buf_mask, buf_value) -> (torch.Tensor, torch.Tensor):\n",
    "        \"\"\"compute the excepted discounted episode return\n",
    "        :int buf_len: the length of ReplayBuffer\n",
    "        :torch.Tensor buf_reward: buf_reward.shape==(buf_len, 1)\n",
    "        :torch.Tensor buf_mask:   buf_mask.shape  ==(buf_len, 1)\n",
    "        :torch.Tensor buf_value:  buf_value.shape ==(buf_len, 1)\n",
    "        :return torch.Tensor buf_r_sum:      buf_r_sum.shape     ==(buf_len, 1)\n",
    "        :return torch.Tensor buf_advantage:  buf_advantage.shape ==(buf_len, 1)\n",
    "        \"\"\"\n",
    "        buf_r_sum = torch.empty(buf_len, dtype=torch.float32, device=self.device)  # reward sum\n",
    "        pre_r_sum = 0  # reward sum of previous step\n",
    "        for i in range(buf_len - 1, -1, -1):\n",
    "            buf_r_sum[i] = buf_reward[i] + buf_mask[i] * pre_r_sum\n",
    "            pre_r_sum = buf_r_sum[i]\n",
    "        buf_advantage = buf_r_sum - (buf_mask * buf_value.squeeze(1))\n",
    "        buf_advantage = (buf_advantage - buf_advantage.mean()) / (buf_advantage.std() + 1e-5)\n",
    "        return buf_r_sum, buf_advantage\n",
    "\n",
    "    def compute_reward_gae(self, buf_len, buf_reward, buf_mask, buf_value) -> (torch.Tensor, torch.Tensor):\n",
    "        \"\"\"compute the excepted discounted episode return\n",
    "        :int buf_len: the length of ReplayBuffer\n",
    "        :torch.Tensor buf_reward: buf_reward.shape==(buf_len, 1)\n",
    "        :torch.Tensor buf_mask:   buf_mask.shape  ==(buf_len, 1)\n",
    "        :torch.Tensor buf_value:  buf_value.shape ==(buf_len, 1)\n",
    "        :return torch.Tensor buf_r_sum:      buf_r_sum.shape     ==(buf_len, 1)\n",
    "        :return torch.Tensor buf_advantage:  buf_advantage.shape ==(buf_len, 1)\n",
    "        \"\"\"\n",
    "        buf_r_sum = torch.empty(buf_len, dtype=torch.float32, device=self.device)  # old policy value\n",
    "        buf_advantage = torch.empty(buf_len, dtype=torch.float32, device=self.device)  # advantage value\n",
    "\n",
    "        pre_r_sum = 0  # reward sum of previous step\n",
    "        pre_advantage = 0  # advantage value of previous step\n",
    "        for i in range(buf_len - 1, -1, -1):\n",
    "            buf_r_sum[i] = buf_reward[i] + buf_mask[i] * pre_r_sum\n",
    "            pre_r_sum = buf_r_sum[i]\n",
    "\n",
    "            buf_advantage[i] = buf_reward[i] + buf_mask[i] * (pre_advantage - buf_value[i])\n",
    "            pre_advantage = buf_value[i] + buf_advantage[i] * self.lambda_gae_adv\n",
    "\n",
    "        buf_advantage = (buf_advantage - buf_advantage.mean()) / (buf_advantage.std() + 1e-5)\n",
    "        return buf_r_sum, buf_advantage\n",
    "\n",
    "\n",
    "class AgentSharedPPO(AgentPPO):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.clip = 0.25  # ratio.clamp(1 - clip, 1 + clip)\n",
    "        self.lambda_entropy = 0.01  # could be 0.02\n",
    "        self.lambda_gae_adv = 0.98  # could be 0.95~0.99, GAE (Generalized Advantage Estimation. ICLR.2016.)\n",
    "        self.obj_c = (-np.log(0.5)) ** 0.5  # for reliable_lambda\n",
    "\n",
    "    def init(self, net_dim, state_dim, action_dim, if_per=False):\n",
    "        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "        self.act = SharedPPO(state_dim, action_dim, net_dim).to(self.device)\n",
    "\n",
    "        self.optimizer = torch.optim.Adam([\n",
    "            {'params': self.act.enc_s.parameters(), 'lr': self.learning_rate * 0.9},\n",
    "            {'params': self.act.dec_a.parameters(), },\n",
    "            {'params': self.act.a_std_log, },\n",
    "            {'params': self.act.dec_q1.parameters(), },\n",
    "            {'params': self.act.dec_q2.parameters(), },\n",
    "        ], lr=self.learning_rate)\n",
    "        self.criterion = torch.nn.SmoothL1Loss()\n",
    "        assert if_per is False  # on-policy don't need PER\n",
    "\n",
    "    def update_net(self, buffer, _target_step, batch_size, repeat_times=4) -> (float, float):  # old version\n",
    "        buffer.update_now_len_before_sample()\n",
    "        buf_len = buffer.now_len  # assert buf_len >= _target_step\n",
    "\n",
    "        '''Trajectory using Generalized Advantage Estimation (GAE)'''\n",
    "        with torch.no_grad():\n",
    "            buf_reward, buf_mask, buf_action, buf_noise, buf_state = buffer.sample_all()\n",
    "\n",
    "            bs = 2 ** 10  # set a smaller 'bs: batch size' when out of GPU memory.\n",
    "            buf_value = torch.cat([self.cri(buf_state[i:i + bs]) for i in range(0, buf_state.size(0), bs)], dim=0)\n",
    "            buf_logprob = -(buf_noise.pow(2).__mul__(0.5) + self.act.a_std_log + self.act.sqrt_2pi_log).sum(1)\n",
    "\n",
    "            buf_r_sum = torch.empty(buf_len, dtype=torch.float32, device=self.device)  # old policy value\n",
    "            buf_advantage = torch.empty(buf_len, dtype=torch.float32, device=self.device)  # advantage value\n",
    "\n",
    "            pre_r_sum = 0  # reward sum of previous step\n",
    "            pre_advantage = 0  # advantage value of previous step\n",
    "            for i in range(buf_len - 1, -1, -1):\n",
    "                buf_r_sum[i] = buf_reward[i] + buf_mask[i] * pre_r_sum\n",
    "                pre_r_sum = buf_r_sum[i]\n",
    "\n",
    "                buf_advantage[i] = buf_reward[i] + buf_mask[i] * (pre_advantage - buf_value[i])\n",
    "                pre_advantage = buf_value[i] + buf_advantage[i] * self.lambda_gae_adv\n",
    "\n",
    "            buf_advantage = (buf_advantage - buf_advantage.mean()) / (buf_advantage.std() + 1e-5)\n",
    "            del buf_reward, buf_mask, buf_noise\n",
    "\n",
    "        '''PPO: Clipped Surrogate objective of Trust Region'''\n",
    "        for _ in range(int(repeat_times * buf_len / batch_size)):\n",
    "            indices = torch.randint(buf_len, size=(batch_size,), device=self.device)\n",
    "\n",
    "            state = buf_state[indices]\n",
    "            action = buf_action[indices]\n",
    "            advantage = buf_advantage[indices]\n",
    "            old_value = buf_r_sum[indices]\n",
    "            old_logprob = buf_logprob[indices]\n",
    "\n",
    "            new_logprob = self.act.compute_logprob(state, action)  # it is obj_actor\n",
    "            ratio = (new_logprob - old_logprob).exp()\n",
    "            obj_surrogate1 = advantage * ratio\n",
    "            obj_surrogate2 = advantage * ratio.clamp(1 - self.clip, 1 + self.clip)\n",
    "            obj_surrogate = -torch.min(obj_surrogate1, obj_surrogate2).mean()\n",
    "            obj_entropy = (new_logprob.exp() * new_logprob).mean()  # policy entropy\n",
    "            obj_actor = obj_surrogate + obj_entropy * self.lambda_entropy\n",
    "\n",
    "            new_value = self.cri(state).squeeze(1)\n",
    "            obj_critic = self.criterion(new_value, old_value)\n",
    "            self.obj_c = 0.995 * self.obj_c + 0.005 * obj_critic.item()  # for reliable_lambda\n",
    "            reliable_lambda = np.exp(-self.obj_c ** 2)  # for reliable_lambda\n",
    "\n",
    "            obj_united = obj_actor * reliable_lambda + obj_critic / (old_value.std() + 1e-5)\n",
    "            self.optimizer.zero_grad()\n",
    "            obj_united.backward()\n",
    "            self.optimizer.step()\n",
    "\n",
    "        return obj_actor.item(), self.obj_c\n",
    "\n",
    "\n",
    "'''Utils'''\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import numpy.random as rd\n",
    "import gym\n",
    "\n",
    "gym.logger.set_level(40)  # Block warning: 'WARN: Box bound precision lowered by casting to float32'\n",
    "\n",
    "\"\"\"[ElegantRL](https://github.com/AI4Finance-LLC/ElegantRL)\"\"\"\n",
    "\n",
    "\n",
    "class PreprocessEnv(gym.Wrapper):  # environment wrapper # todo 2021-03-17\n",
    "    def __init__(self, env, if_print=True, data_type=np.float32):\n",
    "        \"\"\"Preprocess a standard OpenAI gym environment for RL training.\n",
    "        :param env: a standard OpenAI gym environment, it has env.reset() and env.step()\n",
    "        :param if_print: print the information of environment. Such as env_name, state_dim ...\n",
    "        :param data_type: convert state (sometimes float64) to data_type (float32).\n",
    "        \"\"\"\n",
    "        self.env = gym.make(env) if isinstance(env, str) else env\n",
    "        super(PreprocessEnv, self).__init__(self.env)\n",
    "        self.data_type = data_type\n",
    "\n",
    "        (self.env_name, self.state_dim, self.action_dim, self.action_max, self.max_step,\n",
    "         self.if_discrete, self.target_return\n",
    "         ) = get_gym_env_info(self.env, if_print)\n",
    "\n",
    "        state_avg, state_std = get_avg_std__for_state_norm(self.env_name)\n",
    "        if state_avg is not None:\n",
    "            self.neg_state_avg = -state_avg\n",
    "            self.div_state_std = 1 / (state_std + 1e-4)\n",
    "\n",
    "            self.reset = self.reset_norm\n",
    "            self.step = self.step_norm\n",
    "        else:\n",
    "            self.reset = self.reset_type\n",
    "            self.step = self.step_type\n",
    "\n",
    "    def reset_type(self) -> np.ndarray:\n",
    "        \"\"\" state = env.reset()\n",
    "        convert the data type of state from float64 to float32\n",
    "        :return array state: state.shape==(state_dim, )\n",
    "        \"\"\"\n",
    "        state = self.env.reset()\n",
    "        return state.astype(self.data_type)\n",
    "\n",
    "    def reset_norm(self) -> np.ndarray:\n",
    "        \"\"\" state = env.reset()\n",
    "        convert the data type of state from float64 to float32\n",
    "        do normalization on state\n",
    "        :return array state: state.shape==(state_dim, )\n",
    "        \"\"\"\n",
    "        state = self.env.reset()\n",
    "        state = (state + self.neg_state_avg) * self.div_state_std\n",
    "        return state.astype(self.data_type)\n",
    "\n",
    "    def step_type(self, action) -> (np.ndarray, float, bool, dict):\n",
    "        \"\"\" next_state, reward, done = env.step(action)\n",
    "        convert the data type of state from float64 to float32,\n",
    "        adjust action range to (-action_max, +action_max)\n",
    "        :return array state:  state.shape==(state_dim, )\n",
    "        :return float reward: reward of one step\n",
    "        :return bool  done  : the terminal of an training episode\n",
    "        :return dict  info  : the information save in a dict. OpenAI gym standard. Send a `None` is OK\n",
    "        \"\"\"\n",
    "        state, reward, done, info = self.env.step(action * self.action_max)\n",
    "        return state.astype(self.data_type), reward, done, info\n",
    "\n",
    "    def step_norm(self, action) -> (np.ndarray, float, bool, dict):\n",
    "        \"\"\" next_state, reward, done = env.step(action)\n",
    "        convert the data type of state from float64 to float32,\n",
    "        adjust action range to (-action_max, +action_max)\n",
    "        do normalization on state\n",
    "        :return array state:  state.shape==(state_dim, )\n",
    "        :return float reward: reward of one step\n",
    "        :return bool  done  : the terminal of an training episode\n",
    "        :return dict  info  : the information save in a dict. OpenAI gym standard. Send a `None` is OK\n",
    "        \"\"\"\n",
    "        state, reward, done, info = self.env.step(action * self.action_max)\n",
    "        state = (state + self.neg_state_avg) * self.div_state_std\n",
    "        return state.astype(self.data_type), reward, done, info\n",
    "\n",
    "\n",
    "def get_avg_std__for_state_norm(env_name) -> (np.ndarray, np.ndarray):\n",
    "    \"\"\"return the state normalization data: neg_avg and div_std\n",
    "    ReplayBuffer.print_state_norm() will print `neg_avg` and `div_std`\n",
    "    You can save these array to here. And PreprocessEnv will load them automatically.\n",
    "    eg. `state = (state + self.neg_state_avg) * self.div_state_std` in `PreprocessEnv.step_norm()`\n",
    "    neg_avg = -states.mean()\n",
    "    div_std = 1/(states.std()+1e-5) or 6/(states.max()-states.min())\n",
    "    :str env_name: the name of environment that helps to find neg_avg and div_std\n",
    "    :return array avg: neg_avg.shape=(state_dim)\n",
    "    :return array std: div_std.shape=(state_dim)\n",
    "    \"\"\"\n",
    "    avg = None\n",
    "    std = None\n",
    "    if env_name == 'LunarLanderContinuous-v2':\n",
    "        avg = np.array([1.65470898e-02, -1.29684399e-01, 4.26883133e-03, -3.42124557e-02,\n",
    "                        -7.39076972e-03, -7.67103031e-04, 1.12640885e+00, 1.12409466e+00])\n",
    "        std = np.array([0.15094465, 0.29366297, 0.23490797, 0.25931464, 0.21603736,\n",
    "                        0.25886878, 0.277233, 0.27771219])\n",
    "    elif env_name == \"BipedalWalker-v3\":\n",
    "        avg = np.array([1.42211734e-01, -2.74547996e-03, 1.65104509e-01, -1.33418152e-02,\n",
    "                        -2.43243194e-01, -1.73886203e-02, 4.24114229e-02, -6.57800099e-02,\n",
    "                        4.53460692e-01, 6.08022244e-01, -8.64884810e-04, -2.08789053e-01,\n",
    "                        -2.92092949e-02, 5.04791247e-01, 3.33571745e-01, 3.37325723e-01,\n",
    "                        3.49106580e-01, 3.70363115e-01, 4.04074671e-01, 4.55838055e-01,\n",
    "                        5.36685407e-01, 6.70771701e-01, 8.80356865e-01, 9.97987386e-01])\n",
    "        std = np.array([0.84419678, 0.06317835, 0.16532085, 0.09356959, 0.486594,\n",
    "                        0.55477525, 0.44076614, 0.85030824, 0.29159821, 0.48093035,\n",
    "                        0.50323634, 0.48110776, 0.69684234, 0.29161077, 0.06962932,\n",
    "                        0.0705558, 0.07322677, 0.07793258, 0.08624322, 0.09846895,\n",
    "                        0.11752805, 0.14116005, 0.13839757, 0.07760469])\n",
    "    elif env_name == 'ReacherBulletEnv-v0':\n",
    "        avg = np.array([0.03149641, 0.0485873, -0.04949671, -0.06938662, -0.14157104,\n",
    "                        0.02433294, -0.09097818, 0.4405931, 0.10299437], dtype=np.float32)\n",
    "        std = np.array([0.12277275, 0.1347579, 0.14567468, 0.14747661, 0.51311225,\n",
    "                        0.5199606, 0.2710207, 0.48395795, 0.40876198], dtype=np.float32)\n",
    "    elif env_name == 'AntBulletEnv-v0':\n",
    "        avg = np.array([-1.4400886e-01, -4.5074993e-01, 8.5741436e-01, 4.4249415e-01,\n",
    "                        -3.1593361e-01, -3.4174921e-03, -6.1666980e-02, -4.3752361e-03,\n",
    "                        -8.9226037e-02, 2.5108769e-03, -4.8667483e-02, 7.4835382e-03,\n",
    "                        3.6160579e-01, 2.6877613e-03, 4.7474738e-02, -5.0628246e-03,\n",
    "                        -2.5761038e-01, 5.9789192e-04, -2.1119279e-01, -6.6801407e-03,\n",
    "                        2.5196713e-01, 1.6556121e-03, 1.0365561e-01, 1.0219718e-02,\n",
    "                        5.8209229e-01, 7.7563477e-01, 4.8815918e-01, 4.2498779e-01],\n",
    "                       dtype=np.float32)\n",
    "        std = np.array([0.04128463, 0.19463477, 0.15422264, 0.16463493, 0.16640785,\n",
    "                        0.08266512, 0.10606721, 0.07636797, 0.7229637, 0.52585346,\n",
    "                        0.42947173, 0.20228386, 0.44787514, 0.33257666, 0.6440182,\n",
    "                        0.38659114, 0.6644085, 0.5352245, 0.45194066, 0.20750992,\n",
    "                        0.4599643, 0.3846344, 0.651452, 0.39733195, 0.49320385,\n",
    "                        0.41713253, 0.49984455, 0.4943505], dtype=np.float32)\n",
    "    elif env_name == 'HumanoidBulletEnv-v0':\n",
    "        avg = np.array([3.01311314e-01, 3.94672394e-01, 5.94191194e-01, 9.21207070e-02,\n",
    "                        8.33693743e-02, -2.25237925e-02, -1.47895187e-01, 1.78729534e-01,\n",
    "                        6.70446038e-01, 2.97898590e-03, -2.20266372e-01, -1.77605520e-03,\n",
    "                        3.16219926e-02, 5.31213591e-05, 9.07107890e-02, 4.43269382e-04,\n",
    "                        1.03915334e-01, 1.13022688e-04, 8.76481831e-01, -1.43467057e-02,\n",
    "                        7.09028721e-01, -1.55864991e-02, 5.70354581e-01, -3.15685221e-03,\n",
    "                        2.38433480e-01, -1.29739009e-03, 9.75960970e-01, -8.02631397e-03,\n",
    "                        7.48393297e-01, -1.47348447e-02, 4.22917247e-01, -3.47030745e-03,\n",
    "                        7.14308694e-02, -3.49211530e-03, 4.82423425e-01, -7.32147601e-05,\n",
    "                        -5.24461150e-01, -2.18287203e-03, -1.47674218e-01, -3.43166990e-04,\n",
    "                        8.65057111e-02, -2.88956566e-03, 6.23931885e-01, 5.93078613e-01],\n",
    "                       dtype=np.float32)\n",
    "        std = np.array([0.06094389, 0.48514748, 0.4642684, 0.11566383, 0.12077816,\n",
    "                        0.1104386, 0.3986176, 0.3980264, 0.35443318, 0.3776695,\n",
    "                        0.49051976, 0.30215684, 0.615806, 0.40625623, 0.6169094,\n",
    "                        0.33979985, 0.449475, 0.54336107, 0.2420163, 0.37398043,\n",
    "                        0.22486377, 0.4408496, 0.585778, 0.31608477, 0.4161703,\n",
    "                        0.4983718, 0.07819878, 0.29232258, 0.19291587, 0.39967823,\n",
    "                        0.45776755, 0.19698475, 0.48533973, 0.2996624, 0.59454864,\n",
    "                        0.6142501, 0.38873306, 0.19519839, 0.47335255, 0.29171264,\n",
    "                        0.690289, 0.61651593, 0.48313695, 0.4909233], dtype=np.float32)\n",
    "    # elif env_name == 'MinitaurBulletEnv-v0': # need check\n",
    "    #     # avg = np.array([0.90172989, 1.54730119, 1.24560906, 1.97365306, 1.9413892,\n",
    "    #     #                 1.03866835, 1.69646277, 1.18655352, -0.45842347, 0.17845232,\n",
    "    #     #                 0.38784456, 0.58572877, 0.91414561, -0.45410697, 0.7591031,\n",
    "    #     #                 -0.07008998, 3.43842258, 0.61032482, 0.86689961, -0.33910894,\n",
    "    #     #                 0.47030415, 4.5623528, -2.39108079, 3.03559422, -0.36328256,\n",
    "    #     #                 -0.20753499, -0.47758384, 0.86756409])\n",
    "    #     # std = np.array([0.34192648, 0.51169916, 0.39370621, 0.55568461, 0.46910769,\n",
    "    #     #                 0.28387504, 0.51807949, 0.37723445, 13.16686185, 17.51240024,\n",
    "    #     #                 14.80264211, 16.60461412, 15.72930229, 11.38926597, 15.40598346,\n",
    "    #     #                 13.03124941, 2.47718145, 2.55088804, 2.35964651, 2.51025567,\n",
    "    #     #                 2.66379017, 2.37224904, 2.55892521, 2.41716885, 0.07529733,\n",
    "    #     #                 0.05903034, 0.1314812, 0.0221248])\n",
    "    # elif env_name == \"BipedalWalkerHardcore-v3\": # need check\n",
    "    #     avg = np.array([-3.6378160e-02, -2.5788052e-03, 3.4413573e-01, -8.4189959e-03,\n",
    "    #                     -9.1864385e-02, 3.2804706e-04, -6.4693891e-02, -9.8939031e-02,\n",
    "    #                     3.5180664e-01, 6.8103075e-01, 2.2930240e-03, -4.5893672e-01,\n",
    "    #                     -7.6047562e-02, 4.6414185e-01, 3.9363885e-01, 3.9603019e-01,\n",
    "    #                     4.0758255e-01, 4.3053803e-01, 4.6186063e-01, 5.0293463e-01,\n",
    "    #                     5.7822973e-01, 6.9820738e-01, 8.9829963e-01, 9.8080903e-01])\n",
    "    #     std = np.array([0.5771428, 0.05302362, 0.18906464, 0.10137994, 0.41284004,\n",
    "    #                     0.68852615, 0.43710527, 0.87153363, 0.3210142, 0.36864948,\n",
    "    #                     0.6926624, 0.38297284, 0.76805115, 0.33138904, 0.09618598,\n",
    "    #                     0.09843876, 0.10035378, 0.11045089, 0.11910835, 0.13400233,\n",
    "    #                     0.15718603, 0.17106676, 0.14363566, 0.10100251])\n",
    "    return avg, std\n",
    "\n",
    "\n",
    "def get_gym_env_info(env, if_print) -> (str, int, int, int, int, bool, float):\n",
    "    \"\"\"get information of a standard OpenAI gym env.\n",
    "    The DRL algorithm AgentXXX need these env information for building networks and training.\n",
    "    env_name: the environment name, such as XxxXxx-v0\n",
    "    state_dim: the dimension of state\n",
    "    action_dim: the dimension of continuous action; Or the number of discrete action\n",
    "    action_max: the max action of continuous action; action_max == 1 when it is discrete action space\n",
    "    if_discrete: Is this env a discrete action space?\n",
    "    target_return: the target episode return, if agent reach this score, then it pass this game (env).\n",
    "    max_step: the steps in an episode. (from env.reset to done). It breaks an episode when it reach max_step\n",
    "    :env: a standard OpenAI gym environment, it has env.reset() and env.step()\n",
    "    :bool if_print: print the information of environment. Such as env_name, state_dim ...\n",
    "    \"\"\"\n",
    "    gym.logger.set_level(40)  # Block warning: 'WARN: Box bound precision lowered by casting to float32'\n",
    "    assert isinstance(env, gym.Env)\n",
    "\n",
    "    env_name = env.unwrapped.spec.id\n",
    "\n",
    "    state_shape = env.observation_space.shape\n",
    "    state_dim = state_shape[0] if len(state_shape) == 1 else state_shape  # sometimes state_dim is a list\n",
    "\n",
    "    target_return = getattr(env, 'target_return', None)\n",
    "    target_return_default = getattr(env.spec, 'reward_threshold', None)\n",
    "    if target_return is None:\n",
    "        target_return = target_return_default\n",
    "    if target_return is None:\n",
    "        target_return = 2 ** 16\n",
    "\n",
    "    max_step = getattr(env, 'max_step', None)\n",
    "    max_step_default = getattr(env, '_max_episode_steps', None)\n",
    "    if max_step is None:\n",
    "        max_step = max_step_default\n",
    "    if max_step is None:\n",
    "        max_step = 2 ** 10\n",
    "\n",
    "    if_discrete = isinstance(env.action_space, gym.spaces.Discrete)\n",
    "    if if_discrete:  # make sure it is discrete action space\n",
    "        action_dim = env.action_space.n\n",
    "        action_max = int(1)\n",
    "    elif isinstance(env.action_space, gym.spaces.Box):  # make sure it is continuous action space\n",
    "        action_dim = env.action_space.shape[0]\n",
    "        action_max = float(env.action_space.high[0])\n",
    "        assert not any(env.action_space.high + env.action_space.low)\n",
    "    else:\n",
    "        raise RuntimeError('| Please set these value manually: if_discrete=bool, action_dim=int, action_max=1.0')\n",
    "\n",
    "    print(f\"\\n| env_name:  {env_name}, action space if_discrete: {if_discrete}\"\n",
    "          f\"\\n| state_dim: {state_dim:4}, action_dim: {action_dim}, action_max: {action_max}\"\n",
    "          f\"\\n| max_step:  {max_step:4}, target_return: {target_return}\") if if_print else None\n",
    "    return env_name, state_dim, action_dim, action_max, max_step, if_discrete, target_return\n",
    "\n",
    "\n",
    "\"\"\"Custom environment: Finance RL, Github AI4Finance-LLC\"\"\"\n",
    "\n",
    "\n",
    "class FinanceStockEnv:  # 2021-02-02\n",
    "    \"\"\"FinRL\n",
    "    Paper: A Deep Reinforcement Learning Library for Automated Stock Trading in Quantitative Finance\n",
    "           https://arxiv.org/abs/2011.09607 NeurIPS 2020: Deep RL Workshop.\n",
    "    Source: Github https://github.com/AI4Finance-LLC/FinRL-Library\n",
    "    Modify: Github Yonv1943 ElegantRL\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, initial_account=1e6, max_stock=1e2, transaction_fee_percent=1e-3, if_train=True,\n",
    "                 train_beg=0, train_len=1024):\n",
    "        self.stock_dim = 30\n",
    "        self.initial_account = initial_account\n",
    "        self.transaction_fee_percent = transaction_fee_percent\n",
    "        self.max_stock = max_stock\n",
    "\n",
    "        ary = self.load_training_data_for_multi_stock(data_path='./FinanceStock.npy')\n",
    "        assert ary.shape == (1699, 5 * 30)  # ary: (date, item*stock_dim), item: (adjcp, macd, rsi, cci, adx)\n",
    "        assert train_beg < train_len\n",
    "        assert train_len < ary.shape[0]  # ary.shape[0] == 1699\n",
    "        self.ary_train = ary[:train_len]\n",
    "        self.ary_valid = ary[train_len:]\n",
    "        self.ary = self.ary_train if if_train else self.ary_valid\n",
    "\n",
    "        # reset\n",
    "        self.day = 0\n",
    "        self.initial_account__reset = self.initial_account\n",
    "        self.account = self.initial_account__reset\n",
    "        self.day_npy = self.ary[self.day]\n",
    "        self.stocks = np.zeros(self.stock_dim, dtype=np.float32)  # multi-stack\n",
    "\n",
    "        self.total_asset = self.account + (self.day_npy[:self.stock_dim] * self.stocks).sum()\n",
    "        self.episode_return = 0.0  # Compatibility for ElegantRL 2020-12-21\n",
    "        self.gamma_return = 0.0\n",
    "\n",
    "        '''env information'''\n",
    "        self.env_name = 'FinanceStock-v2'\n",
    "        self.state_dim = 1 + (5 + 1) * self.stock_dim\n",
    "        self.action_dim = self.stock_dim\n",
    "        self.if_discrete = False\n",
    "        self.target_return = 1.25  # convergence 1.5\n",
    "        self.max_step = self.ary.shape[0]\n",
    "\n",
    "    def reset(self) -> np.ndarray:\n",
    "        self.initial_account__reset = self.initial_account * rd.uniform(0.9, 1.1)  # reset()\n",
    "        self.account = self.initial_account__reset\n",
    "        self.stocks = np.zeros(self.stock_dim, dtype=np.float32)\n",
    "        self.total_asset = self.account + (self.day_npy[:self.stock_dim] * self.stocks).sum()\n",
    "        # total_asset = account + (adjcp * stocks).sum()\n",
    "\n",
    "        self.day = 0\n",
    "        self.day_npy = self.ary[self.day]\n",
    "        self.day += 1\n",
    "\n",
    "        state = np.hstack((self.account * 2 ** -16,\n",
    "                           self.day_npy * 2 ** -8,\n",
    "                           self.stocks * 2 ** -12,), ).astype(np.float32)\n",
    "        return state\n",
    "\n",
    "    def step(self, action) -> (np.ndarray, float, bool, None):\n",
    "        action = action * self.max_stock\n",
    "\n",
    "        \"\"\"bug or sell stock\"\"\"\n",
    "        for index in range(self.stock_dim):\n",
    "            stock_action = action[index]\n",
    "            adj = self.day_npy[index]\n",
    "            if stock_action > 0:  # buy_stock\n",
    "                available_amount = self.account // adj\n",
    "                delta_stock = min(available_amount, stock_action)\n",
    "                self.account -= adj * delta_stock * (1 + self.transaction_fee_percent)\n",
    "                self.stocks[index] += delta_stock\n",
    "            elif self.stocks[index] > 0:  # sell_stock\n",
    "                delta_stock = min(-stock_action, self.stocks[index])\n",
    "                self.account += adj * delta_stock * (1 - self.transaction_fee_percent)\n",
    "                self.stocks[index] -= delta_stock\n",
    "\n",
    "        \"\"\"update day\"\"\"\n",
    "        self.day_npy = self.ary[self.day]\n",
    "        self.day += 1\n",
    "        done = self.day == self.max_step  # 2020-12-21\n",
    "\n",
    "        state = np.hstack((self.account * 2 ** -16,\n",
    "                           self.day_npy * 2 ** -8,\n",
    "                           self.stocks * 2 ** -12,), ).astype(np.float32)\n",
    "\n",
    "        next_total_asset = self.account + (self.day_npy[:self.stock_dim] * self.stocks).sum()\n",
    "        reward = (next_total_asset - self.total_asset) * 2 ** -16  # notice scaling!\n",
    "        self.total_asset = next_total_asset\n",
    "\n",
    "        self.gamma_return = self.gamma_return * 0.99 + reward  # notice: gamma_r seems good? Yes\n",
    "        if done:\n",
    "            reward += self.gamma_return\n",
    "            self.gamma_return = 0.0  # env.reset()\n",
    "\n",
    "            # cumulative_return_rate\n",
    "            self.episode_return = next_total_asset / self.initial_account\n",
    "\n",
    "        return state, reward, done, None\n",
    "\n",
    "    @staticmethod\n",
    "    def load_training_data_for_multi_stock(data_path='./FinanceStock.npy'):  # need more independent\n",
    "        if os.path.exists(data_path):\n",
    "            data_ary = np.load(data_path).astype(np.float32)\n",
    "            assert data_ary.shape[1] == 5 * 30\n",
    "            return data_ary\n",
    "        else:\n",
    "            raise RuntimeError(\n",
    "                f'| Download and put it into: {data_path}\\n for FinanceStockEnv()'\n",
    "                f'| https://github.com/Yonv1943/ElegantRL/blob/master/FinanceMultiStock.npy'\n",
    "                f'| Or you can use the following code to generate it from a csv file.')\n",
    "\n",
    "        # from preprocessing.preprocessors import pd, data_split, preprocess_data, add_turbulence\n",
    "        #\n",
    "        # # the following is same as part of run_model()\n",
    "        # preprocessed_path = \"done_data.csv\"\n",
    "        # if if_load and os.path.exists(preprocessed_path):\n",
    "        #     data = pd.read_csv(preprocessed_path, index_col=0)\n",
    "        # else:\n",
    "        #     data = preprocess_data()\n",
    "        #     data = add_turbulence(data)\n",
    "        #     data.to_csv(preprocessed_path)\n",
    "        #\n",
    "        # df = data\n",
    "        # rebalance_window = 63\n",
    "        # validation_window = 63\n",
    "        # i = rebalance_window + validation_window\n",
    "        #\n",
    "        # unique_trade_date = data[(data.datadate > 20151001) & (data.datadate <= 20200707)].datadate.unique()\n",
    "        # train__df = data_split(df, start=20090000, end=unique_trade_date[i - rebalance_window - validation_window])\n",
    "        # # print(train__df) # df: DataFrame of Pandas\n",
    "        #\n",
    "        # train_ary = train__df.to_numpy().reshape((-1, 30, 12))\n",
    "        # '''state_dim = 1 + 6 * stock_dim, stock_dim=30\n",
    "        # n   item    index\n",
    "        # 1   ACCOUNT -\n",
    "        # 30  adjcp   2\n",
    "        # 30  stock   -\n",
    "        # 30  macd    7\n",
    "        # 30  rsi     8\n",
    "        # 30  cci     9\n",
    "        # 30  adx     10\n",
    "        # '''\n",
    "        # data_ary = np.empty((train_ary.shape[0], 5, 30), dtype=np.float32)\n",
    "        # data_ary[:, 0] = train_ary[:, :, 2]  # adjcp\n",
    "        # data_ary[:, 1] = train_ary[:, :, 7]  # macd\n",
    "        # data_ary[:, 2] = train_ary[:, :, 8]  # rsi\n",
    "        # data_ary[:, 3] = train_ary[:, :, 9]  # cci\n",
    "        # data_ary[:, 4] = train_ary[:, :, 10]  # adx\n",
    "        #\n",
    "        # data_ary = data_ary.reshape((-1, 5 * 30))\n",
    "        #\n",
    "        # os.makedirs(data_path[:data_path.rfind('/')])\n",
    "        # np.save(data_path, data_ary.astype(np.float16))  # save as float16 (0.5 MB), float32 (1.0 MB)\n",
    "        # print('| FinanceStockEnv(): save in:', data_path)\n",
    "        # return data_ary\n",
    "\n",
    "    def draw_cumulative_return(self, args, torch) -> list:\n",
    "        state_dim = self.state_dim\n",
    "        action_dim = self.action_dim\n",
    "\n",
    "        agent_rl = args.agent\n",
    "        net_dim = args.net_dim\n",
    "        cwd = args.cwd\n",
    "\n",
    "        agent = agent_rl(net_dim, state_dim, action_dim)  # build AgentRL\n",
    "        act = agent.act\n",
    "        device = agent.device\n",
    "\n",
    "        state = self.reset()\n",
    "        episode_returns = list()  # the cumulative_return / initial_account\n",
    "        with torch.no_grad():\n",
    "            for i in range(self.max_step):\n",
    "                s_tensor = torch.as_tensor((state,), device=device)\n",
    "                a_tensor = act(s_tensor)\n",
    "                action = a_tensor.cpu().numpy()[0]  # not need detach(), because with torch.no_grad() outside\n",
    "                state, reward, done, _ = self.step(action)\n",
    "\n",
    "                episode_return = (self.account + (self.day_npy[:self.stock_dim] * self.stocks).sum()\n",
    "                                  ) / self.initial_account__reset\n",
    "                episode_returns.append(episode_return)\n",
    "                if done:\n",
    "                    break\n",
    "\n",
    "        import matplotlib.pyplot as plt\n",
    "        plt.plot(episode_returns)\n",
    "        plt.grid()\n",
    "        plt.title('cumulative return')\n",
    "        plt.xlabel('day')\n",
    "        plt.xlabel('multiple of initial_account')\n",
    "        plt.savefig(f'{cwd}/cumulative_return.jpg')\n",
    "        return episode_returns\n",
    "\n",
    "\n",
    "\"\"\"Custom environment: Fix Env\"\"\"\n",
    "\n",
    "\n",
    "def fix_car_racing_env(env, frame_num=3, action_num=3) -> gym.Wrapper:  # 2020-12-12\n",
    "    setattr(env, 'old_step', env.step)  # env.old_step = env.step\n",
    "    setattr(env, 'env_name', 'CarRacing-Fix')\n",
    "    setattr(env, 'state_dim', (frame_num, 96, 96))\n",
    "    setattr(env, 'action_dim', 3)\n",
    "    setattr(env, 'if_discrete', False)\n",
    "    setattr(env, 'target_return', 700)  # 900 in default\n",
    "\n",
    "    setattr(env, 'state_stack', None)  # env.state_stack = None\n",
    "    setattr(env, 'avg_reward', 0)  # env.avg_reward = 0\n",
    "    \"\"\" cancel the print() in environment\n",
    "    comment 'car_racing.py' line 233-234: print('Track generation ...\n",
    "    comment 'car_racing.py' line 308-309: print(\"retry to generate track ...\n",
    "    \"\"\"\n",
    "\n",
    "    def rgb2gray(rgb):\n",
    "        # # rgb image -> gray [0, 1]\n",
    "        # gray = np.dot(rgb[..., :], [0.299, 0.587, 0.114]).astype(np.float32)\n",
    "        # if norm:\n",
    "        #     # normalize\n",
    "        #     gray = gray / 128. - 1.\n",
    "        # return gray\n",
    "\n",
    "        state = rgb[:, :, 1]  # show green\n",
    "        state[86:, 24:36] = rgb[86:, 24:36, 2]  # show red\n",
    "        state[86:, 72:] = rgb[86:, 72:, 0]  # show blue\n",
    "        state = (state - 128).astype(np.float32) / 128.\n",
    "        return state\n",
    "\n",
    "    def decorator_step(env_step):\n",
    "        def new_env_step(action):\n",
    "            action = action.copy()\n",
    "            action[1:] = (action[1:] + 1) / 2  # fix action_space.low\n",
    "\n",
    "            reward_sum = 0\n",
    "            done = state = None\n",
    "            try:\n",
    "                for _ in range(action_num):\n",
    "                    state, reward, done, info = env_step(action)\n",
    "                    state = rgb2gray(state)\n",
    "\n",
    "                    if done:\n",
    "                        reward += 100  # don't penalize \"die state\"\n",
    "                    if state.mean() > 192:  # 185.0:  # penalize when outside of road\n",
    "                        reward -= 0.05\n",
    "\n",
    "                    env.avg_reward = env.avg_reward * 0.95 + reward * 0.05\n",
    "                    if env.avg_reward <= -0.1:  # done if car don't move\n",
    "                        done = True\n",
    "\n",
    "                    reward_sum += reward\n",
    "\n",
    "                    if done:\n",
    "                        break\n",
    "            except Exception as error:\n",
    "                print(f\"| CarRacing-v0 Error 'stack underflow'? {error}\")\n",
    "                reward_sum = 0\n",
    "                done = True\n",
    "            env.state_stack.pop(0)\n",
    "            env.state_stack.append(state)\n",
    "\n",
    "            return np.array(env.state_stack).flatten(), reward_sum, done, {}\n",
    "\n",
    "        return new_env_step\n",
    "\n",
    "    env.step = decorator_step(env.step)\n",
    "\n",
    "    def decorator_reset(env_reset):\n",
    "        def new_env_reset():\n",
    "            state = rgb2gray(env_reset())\n",
    "            env.state_stack = [state, ] * frame_num\n",
    "            return np.array(env.state_stack).flatten()\n",
    "\n",
    "        return new_env_reset\n",
    "\n",
    "    env.reset = decorator_reset(env.reset)\n",
    "    return env\n",
    "\n",
    "\n",
    "def render__car_racing():\n",
    "    import gym  # gym of OpenAI is not necessary for ElegantRL (even RL)\n",
    "    gym.logger.set_level(40)  # Block warning: 'WARN: Box bound precision lowered by casting to float32'\n",
    "    env = gym.make('CarRacing-v0')\n",
    "    env = fix_car_racing_env(env)\n",
    "\n",
    "    state_dim = env.state_dim\n",
    "\n",
    "    _state = env.reset()\n",
    "    import cv2\n",
    "    action = np.array((0, 1.0, -1.0))\n",
    "    for i in range(321):\n",
    "        # action = env.action_space.sample()\n",
    "        state, reward, done, _ = env.step(action)\n",
    "        # env.render\n",
    "        show = state.reshape(state_dim)\n",
    "        show = ((show[0] + 1.0) * 128).astype(np.uint8)\n",
    "        cv2.imshow('', show)\n",
    "        cv2.waitKey(1)\n",
    "        if done:\n",
    "            break\n",
    "        # env.render()\n",
    "\n",
    "\n",
    "\"\"\"Utils\"\"\"\n",
    "\n",
    "\n",
    "def get_video_to_watch_gym_render():\n",
    "    import cv2  # pip3 install opencv-python\n",
    "    import gym  # pip3 install gym==0.17 pyglet==1.5.0  # env.render() bug in gym==0.18, pyglet==1.6\n",
    "    import torch\n",
    "\n",
    "    '''choose env'''\n",
    "    # from elegantrl.env import PreprocessEnv\n",
    "    env = PreprocessEnv(env=gym.make('BipedalWalker-v3'))\n",
    "\n",
    "    '''choose algorithm'''\n",
    "    from elegantrl.agent import AgentPPO\n",
    "    agent = AgentPPO()\n",
    "    net_dim = 2 ** 8\n",
    "    cwd = 'AgentPPO/BipedalWalker-v3_2/'\n",
    "    # from elegantrl.agent import AgentModSAC\n",
    "    # agent = AgentModSAC()\n",
    "    # net_dim = 2 ** 7\n",
    "    # cwd = 'AgentModSAC/BipedalWalker-v3_2/'\n",
    "\n",
    "    '''initialize agent'''\n",
    "    state_dim = env.state_dim\n",
    "    action_dim = env.action_dim\n",
    "    agent.init(net_dim, state_dim, action_dim)\n",
    "    agent.save_load_model(cwd=cwd, if_save=False)\n",
    "\n",
    "    '''initialize evaluete and env.render()'''\n",
    "    device = agent.device\n",
    "    save_frame_dir = 'frames'\n",
    "    save_video = 'gym_render.mp4'\n",
    "\n",
    "    os.makedirs(save_frame_dir, exist_ok=True)\n",
    "\n",
    "    state = env.reset()\n",
    "    for i in range(1024):\n",
    "        frame = env.render('rgb_array')\n",
    "        cv2.imwrite(f'{save_frame_dir}/{i:06}.png', frame)\n",
    "        # cv2.imshow('', frame)\n",
    "        # cv2.waitKey(1)\n",
    "\n",
    "        s_tensor = torch.as_tensor((state,), dtype=torch.float32, device=device)\n",
    "        a_tensor = agent.act(s_tensor)\n",
    "        action = a_tensor.detach().cpu().numpy()[0]  # if use 'with torch.no_grad()', then '.detach()' not need.\n",
    "        # action = gym_env.action_space.sample()\n",
    "\n",
    "        next_state, reward, done, _ = env.step(action)\n",
    "\n",
    "        if done:\n",
    "            state = env.reset()\n",
    "        else:\n",
    "            state = next_state\n",
    "    env.close()\n",
    "\n",
    "    '''convert frames png/jpg to video mp4/avi using ffmpeg'''\n",
    "    os.system(f\"| Convert frames to video using ffmpeg. Save in {save_video}\")\n",
    "    os.system(f'ffmpeg -r 60 -f image2 -s 600x400 -i {save_frame_dir}/%06d.png '\n",
    "              f'-crf 25 -vb 20M -pix_fmt yuv420p {save_video}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import numpy.random as rd\n",
    "\n",
    "\"\"\"[ElegantRL](https://github.com/AI4Finance-LLC/ElegantRL)\"\"\"\n",
    "\n",
    "\n",
    "class ReplayBuffer1:\n",
    "    def __init__(self, max_len, state_dim, action_dim, if_on_policy, if_per, if_gpu):\n",
    "        \"\"\"Experience Replay Buffer\n",
    "        save environment transition in a continuous RAM for high performance training\n",
    "        we save trajectory in order and save state and other (action, reward, mask, ...) separately.\n",
    "        `int max_len` the maximum capacity of ReplayBuffer. First In First Out\n",
    "        `int state_dim` the dimension of state\n",
    "        `int action_dim` the dimension of action (action_dim==1 for discrete action)\n",
    "        `bool if_on_policy` on-policy or off-policy\n",
    "        `bool if_gpu` create buffer space on CPU RAM or GPU\n",
    "        `bool if_per` Prioritized Experience Replay for sparse reward\n",
    "        \"\"\"\n",
    "\n",
    "        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "        self.max_len = max_len\n",
    "        self.now_len = 0\n",
    "        self.next_idx = 0\n",
    "        self.if_full = False\n",
    "        self.action_dim = action_dim  # for self.sample_all(\n",
    "        self.if_on_policy = if_on_policy\n",
    "        # self.if_gpu = if_gpu\n",
    "\n",
    "        if if_per:\n",
    "            self.tree = BinarySearchTree(max_len)\n",
    "            self.append_buffer = self.append_buffer_per\n",
    "            self.extend_buffer = self.extend_buffer_per\n",
    "            self.sample_batch = self.sample_batch_per\n",
    "\n",
    "        if if_on_policy:\n",
    "            # self.if_gpu = False\n",
    "            other_dim = 1 + 1 + action_dim * 2\n",
    "        else:\n",
    "            other_dim = 1 + 1 + action_dim\n",
    "\n",
    "        # if self.if_gpu:\n",
    "        #     self.buf_other = torch.empty((max_len, other_dim), dtype=torch.float32, device=self.device)\n",
    "        #     self.buf_state = torch.empty((max_len, state_dim), dtype=torch.float32, device=self.device)\n",
    "        # else:\n",
    "        self.buf_other = np.empty((max_len, other_dim), dtype=np.float32)\n",
    "        self.buf_state = np.empty((max_len, state_dim), dtype=np.float32)\n",
    "\n",
    "    def append_buffer(self, state, other):  # CPU array to CPU array\n",
    "        # if self.if_gpu:\n",
    "        #     state = torch.as_tensor(state, device=self.device)\n",
    "        #     other = torch.as_tensor(other, device=self.device)\n",
    "        self.buf_state[self.next_idx] = state\n",
    "        self.buf_other[self.next_idx] = other\n",
    "\n",
    "        self.next_idx += 1\n",
    "        if self.next_idx >= self.max_len:\n",
    "            self.if_full = True\n",
    "            self.next_idx = 0\n",
    "\n",
    "    def append_buffer_per(self, state, other):  # CPU array to CPU array\n",
    "        # if self.if_gpu:\n",
    "        #     state = torch.as_tensor(state, device=self.device)\n",
    "        #     other = torch.as_tensor(other, device=self.device)\n",
    "        self.buf_state[self.next_idx] = state\n",
    "        self.buf_other[self.next_idx] = other\n",
    "\n",
    "        self.tree.update_id(self.next_idx)\n",
    "\n",
    "        self.next_idx += 1\n",
    "        if self.next_idx >= self.max_len:\n",
    "            self.if_full = True\n",
    "            self.next_idx = 0\n",
    "\n",
    "    def extend_buffer(self, state, other):  # CPU array to CPU array\n",
    "        # if self.if_gpu:\n",
    "        #     state = torch.as_tensor(state, dtype=torch.float32, device=self.device)\n",
    "        #     other = torch.as_tensor(other, dtype=torch.float32, device=self.device)\n",
    "\n",
    "        size = len(other)\n",
    "        next_idx = self.next_idx + size\n",
    "\n",
    "        if next_idx > self.max_len:\n",
    "            if next_idx > self.max_len:\n",
    "                self.buf_state[self.next_idx:self.max_len] = state[:self.max_len - self.next_idx]\n",
    "                self.buf_other[self.next_idx:self.max_len] = other[:self.max_len - self.next_idx]\n",
    "            self.if_full = True\n",
    "            next_idx = next_idx - self.max_len\n",
    "\n",
    "            self.buf_state[0:next_idx] = state[-next_idx:]\n",
    "            self.buf_other[0:next_idx] = other[-next_idx:]\n",
    "        else:\n",
    "            self.buf_state[self.next_idx:next_idx] = state\n",
    "            self.buf_other[self.next_idx:next_idx] = other\n",
    "        self.next_idx = next_idx\n",
    "\n",
    "    def extend_buffer_per(self, state, other):  # CPU array to CPU array\n",
    "        # if self.if_gpu:\n",
    "        #     state = torch.as_tensor(state, dtype=torch.float32, device=self.device)\n",
    "        #     other = torch.as_tensor(other, dtype=torch.float32, device=self.device)\n",
    "\n",
    "        size = len(other)\n",
    "        next_idx = self.next_idx + size\n",
    "\n",
    "        self.tree.update_ids(data_ids=np.arange(self.next_idx, next_idx) % self.max_len)\n",
    "\n",
    "        if next_idx > self.max_len:\n",
    "            if next_idx > self.max_len:\n",
    "                self.buf_state[self.next_idx:self.max_len] = state[:self.max_len - self.next_idx]\n",
    "                self.buf_other[self.next_idx:self.max_len] = other[:self.max_len - self.next_idx]\n",
    "            self.if_full = True\n",
    "            next_idx = next_idx - self.max_len\n",
    "\n",
    "            self.buf_state[0:next_idx] = state[-next_idx:]\n",
    "            self.buf_other[0:next_idx] = other[-next_idx:]\n",
    "        else:\n",
    "            self.buf_state[self.next_idx:next_idx] = state\n",
    "            self.buf_other[self.next_idx:next_idx] = other\n",
    "        self.next_idx = next_idx\n",
    "\n",
    "    def sample_batch(self, batch_size) -> tuple:\n",
    "        \"\"\"randomly sample a batch of data for training\n",
    "        :int batch_size: the number of data in a batch for Stochastic Gradient Descent\n",
    "        :return torch.Tensor reward: reward.shape==(now_len, 1)\n",
    "        :return torch.Tensor mask:   mask.shape  ==(now_len, 1), mask = 0.0 if done else gamma\n",
    "        :return torch.Tensor action: action.shape==(now_len, action_dim)\n",
    "        :return torch.Tensor state:  state.shape ==(now_len, state_dim)\n",
    "        :return torch.Tensor state:  state.shape ==(now_len, state_dim), next state\n",
    "        \"\"\"\n",
    "        indices = rd.randint(self.now_len - 1, size=batch_size)\n",
    "        r_m_a = self.buf_other[indices]\n",
    "        return (r_m_a[:, 0:1],\n",
    "                r_m_a[:, 1:2],\n",
    "                r_m_a[:, 2:],\n",
    "                self.buf_state[indices],\n",
    "                self.buf_state[indices + 1])\n",
    "\n",
    "    def sample_batch_per(self, batch_size) -> tuple:\n",
    "        beg = -self.max_len\n",
    "        end = (self.now_len - self.max_len) if (self.now_len < self.max_len) else None\n",
    "\n",
    "        indices, is_weights = self.tree.get_indices_is_weights(batch_size, beg, end)\n",
    "\n",
    "        r_m_a = self.buf_other[indices]\n",
    "        return (r_m_a[:, 0:1],\n",
    "                r_m_a[:, 1:2],\n",
    "                r_m_a[:, 2:],\n",
    "                self.buf_state[indices],\n",
    "                self.buf_state[indices + 1],\n",
    "                torch.as_tensor(is_weights, dtype=torch.float32, device=self.device))\n",
    "\n",
    "    def sample_all(self) -> tuple:\n",
    "        \"\"\"sample all the data in ReplayBuffer (for on-policy)\n",
    "        :return torch.Tensor reward: reward.shape==(now_len, 1)\n",
    "        :return torch.Tensor mask:   mask.shape  ==(now_len, 1), mask = 0.0 if done else gamma\n",
    "        :return torch.Tensor action: action.shape==(now_len, action_dim)\n",
    "        :return torch.Tensor noise:  noise.shape ==(now_len, action_dim)\n",
    "        :return torch.Tensor state:  state.shape ==(now_len, state_dim)\n",
    "        \"\"\"\n",
    "        all_other = torch.as_tensor(self.buf_other[:self.now_len], device=self.device)\n",
    "        return (all_other[:, 0],\n",
    "                all_other[:, 1],\n",
    "                all_other[:, 2:2 + self.action_dim],\n",
    "                all_other[:, 2 + self.action_dim:],\n",
    "                torch.as_tensor(self.buf_state[:self.now_len], device=self.device))\n",
    "\n",
    "    def update_now_len_before_sample(self):\n",
    "        \"\"\"update the a pointer `now_len`, which is the current data number of ReplayBuffer\n",
    "        \"\"\"\n",
    "        self.now_len = self.max_len if self.if_full else self.next_idx\n",
    "\n",
    "    def empty_buffer_before_explore(self):\n",
    "        \"\"\"we empty the buffer by set now_len=0. On-policy need to empty buffer before exploration\n",
    "        \"\"\"\n",
    "        self.next_idx = 0\n",
    "        self.now_len = 0\n",
    "        self.if_full = False\n",
    "\n",
    "    def print_state_norm(self, neg_avg=None, div_std=None):  # non-essential\n",
    "        \"\"\"print the state norm information: state_avg, state_std\n",
    "        We don't suggest to use running stat state.\n",
    "        We directly do normalization on state using the historical avg and std\n",
    "        eg. `state = (state + self.neg_state_avg) * self.div_state_std` in `PreprocessEnv.step_norm()`\n",
    "        neg_avg = -states.mean()\n",
    "        div_std = 1/(states.std()+1e-5) or 6/(states.max()-states.min())\n",
    "        :array neg_avg: neg_avg.shape=(state_dim)\n",
    "        :array div_std: div_std.shape=(state_dim)\n",
    "        \"\"\"\n",
    "        max_sample_size = 2 ** 14\n",
    "\n",
    "        '''check if pass'''\n",
    "        state_shape = self.buf_state.shape\n",
    "        if len(state_shape) > 2 or state_shape[1] > 64:\n",
    "            print(f\"| print_state_norm(): state_dim: {state_shape} is too large to print its norm. \")\n",
    "            return None\n",
    "\n",
    "        '''sample state'''\n",
    "        indices = np.arange(self.now_len)\n",
    "        rd.shuffle(indices)\n",
    "        indices = indices[:max_sample_size]  # len(indices) = min(self.now_len, max_sample_size)\n",
    "\n",
    "        batch_state = self.buf_state[indices]\n",
    "\n",
    "        '''compute state norm'''\n",
    "        if isinstance(batch_state, torch.Tensor):\n",
    "            batch_state = batch_state.cpu().data.numpy()\n",
    "        assert isinstance(batch_state, np.ndarray)\n",
    "\n",
    "        if batch_state.shape[1] > 64:\n",
    "            print(f\"| _print_norm(): state_dim: {batch_state.shape[1]:.0f} is too large to print its norm. \")\n",
    "            return None\n",
    "\n",
    "        if np.isnan(batch_state).any():  # 2020-12-12\n",
    "            batch_state = np.nan_to_num(batch_state)  # nan to 0\n",
    "\n",
    "        ary_avg = batch_state.mean(axis=0)\n",
    "        ary_std = batch_state.std(axis=0)\n",
    "        fix_std = ((np.max(batch_state, axis=0) - np.min(batch_state, axis=0)) / 6 + ary_std) / 2\n",
    "\n",
    "        if neg_avg is not None:  # norm transfer\n",
    "            ary_avg = ary_avg - neg_avg / div_std\n",
    "            ary_std = fix_std / div_std\n",
    "\n",
    "        print(f\"print_state_norm: state_avg, state_std (fixed)\")\n",
    "        print(f\"avg = np.{repr(ary_avg).replace('=float32', '=np.float32')}\")\n",
    "        print(f\"std = np.{repr(ary_std).replace('=float32', '=np.float32')}\")\n",
    "\n",
    "    def td_error_update(self, td_error):\n",
    "        self.tree.td_error_update(td_error)\n",
    "\n",
    "\n",
    "class ReplayBuffer:\n",
    "    def __init__(self, max_len, state_dim, action_dim, if_on_policy, if_per, if_gpu):\n",
    "        \"\"\"Experience Replay Buffer\n",
    "        save environment transition in a continuous RAM for high performance training\n",
    "        we save trajectory in order and save state and other (action, reward, mask, ...) separately.\n",
    "        `int max_len` the maximum capacity of ReplayBuffer. First In First Out\n",
    "        `int state_dim` the dimension of state\n",
    "        `int action_dim` the dimension of action (action_dim==1 for discrete action)\n",
    "        `bool if_on_policy` on-policy or off-policy\n",
    "        `bool if_gpu` create buffer space on CPU RAM or GPU\n",
    "        `bool if_per` Prioritized Experience Replay for sparse reward\n",
    "        \"\"\"\n",
    "        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "        self.max_len = max_len\n",
    "        self.now_len = 0\n",
    "        self.next_idx = 0\n",
    "        self.if_full = False\n",
    "        self.action_dim = action_dim  # for self.sample_all(\n",
    "        self.if_on_policy = if_on_policy\n",
    "        self.if_gpu = if_gpu\n",
    "\n",
    "        self.if_per = if_per\n",
    "        if if_per:\n",
    "            self.tree = BinarySearchTree(max_len)\n",
    "\n",
    "        if if_on_policy:\n",
    "            self.if_gpu = False\n",
    "            other_dim = 1 + 1 + action_dim * 2\n",
    "        else:\n",
    "            other_dim = 1 + 1 + action_dim\n",
    "\n",
    "        if self.if_gpu:\n",
    "            self.buf_other = torch.empty((max_len, other_dim), dtype=torch.float32, device=self.device)\n",
    "            self.buf_state = torch.empty((max_len, state_dim), dtype=torch.float32, device=self.device)\n",
    "        else:\n",
    "            self.buf_other = np.empty((max_len, other_dim), dtype=np.float32)\n",
    "            self.buf_state = np.empty((max_len, state_dim), dtype=np.float32)\n",
    "\n",
    "    def append_buffer(self, state, other):  # CPU array to CPU array\n",
    "        if self.if_gpu:\n",
    "            state = torch.as_tensor(state, device=self.device)\n",
    "            other = torch.as_tensor(other, device=self.device)\n",
    "        self.buf_state[self.next_idx] = state\n",
    "        self.buf_other[self.next_idx] = other\n",
    "\n",
    "        if self.if_per:\n",
    "            self.tree.update_id(self.next_idx)\n",
    "\n",
    "        self.next_idx += 1\n",
    "        if self.next_idx >= self.max_len:\n",
    "            self.if_full = True\n",
    "            self.next_idx = 0\n",
    "\n",
    "    def extend_buffer(self, state, other):  # CPU array to CPU array\n",
    "        if self.if_gpu:\n",
    "            state = torch.as_tensor(state, dtype=torch.float32, device=self.device)\n",
    "            other = torch.as_tensor(other, dtype=torch.float32, device=self.device)\n",
    "\n",
    "        size = len(other)\n",
    "        next_idx = self.next_idx + size\n",
    "\n",
    "        if self.if_per:\n",
    "            for data_id in (np.arange(self.next_idx, next_idx) % self.max_len):\n",
    "                self.tree.update_ids(data_id)\n",
    "\n",
    "        if next_idx > self.max_len:\n",
    "            if next_idx > self.max_len:\n",
    "                self.buf_state[self.next_idx:self.max_len] = state[:self.max_len - self.next_idx]\n",
    "                self.buf_other[self.next_idx:self.max_len] = other[:self.max_len - self.next_idx]\n",
    "            self.if_full = True\n",
    "            next_idx = next_idx - self.max_len\n",
    "\n",
    "            self.buf_state[0:next_idx] = state[-next_idx:]\n",
    "            self.buf_other[0:next_idx] = other[-next_idx:]\n",
    "        else:\n",
    "            self.buf_state[self.next_idx:next_idx] = state\n",
    "            self.buf_other[self.next_idx:next_idx] = other\n",
    "        self.next_idx = next_idx\n",
    "\n",
    "    def sample_batch(self, batch_size) -> tuple:\n",
    "        \"\"\"randomly sample a batch of data for training\n",
    "        :int batch_size: the number of data in a batch for Stochastic Gradient Descent\n",
    "        :return torch.Tensor reward: reward.shape==(now_len, 1)\n",
    "        :return torch.Tensor mask:   mask.shape  ==(now_len, 1), mask = 0.0 if done else gamma\n",
    "        :return torch.Tensor action: action.shape==(now_len, action_dim)\n",
    "        :return torch.Tensor state:  state.shape ==(now_len, state_dim)\n",
    "        :return torch.Tensor state:  state.shape ==(now_len, state_dim), next state\n",
    "        \"\"\"\n",
    "        if self.if_per:\n",
    "            beg = -self.max_len\n",
    "            end = (self.now_len - self.max_len) if (self.now_len < self.max_len) else None\n",
    "\n",
    "            indices, is_weights = self.tree.get_indices_is_weights(batch_size, beg, end)\n",
    "\n",
    "            r_m_a = self.buf_other[indices]\n",
    "            return (r_m_a[:, 0:1],\n",
    "                    r_m_a[:, 1:2],\n",
    "                    r_m_a[:, 2:],\n",
    "                    self.buf_state[indices],\n",
    "                    self.buf_state[indices + 1],\n",
    "                    torch.as_tensor(is_weights, dtype=torch.float32, device=self.device))\n",
    "\n",
    "        else:\n",
    "            indices = torch.randint(self.now_len - 1, size=(batch_size,), device=self.device) if self.if_gpu \\\n",
    "                else rd.randint(self.now_len - 1, size=batch_size)\n",
    "            r_m_a = self.buf_other[indices]\n",
    "            return (r_m_a[:, 0:1],\n",
    "                    r_m_a[:, 1:2],\n",
    "                    r_m_a[:, 2:],\n",
    "                    self.buf_state[indices],\n",
    "                    self.buf_state[indices + 1])\n",
    "\n",
    "    def sample_all(self) -> tuple:\n",
    "        \"\"\"sample all the data in ReplayBuffer (for on-policy)\n",
    "        :return torch.Tensor reward: reward.shape==(now_len, 1)\n",
    "        :return torch.Tensor mask:   mask.shape  ==(now_len, 1), mask = 0.0 if done else gamma\n",
    "        :return torch.Tensor action: action.shape==(now_len, action_dim)\n",
    "        :return torch.Tensor noise:  noise.shape ==(now_len, action_dim)\n",
    "        :return torch.Tensor state:  state.shape ==(now_len, state_dim)\n",
    "        \"\"\"\n",
    "        all_other = torch.as_tensor(self.buf_other[:self.now_len], device=self.device)\n",
    "        return (all_other[:, 0],\n",
    "                all_other[:, 1],\n",
    "                all_other[:, 2:2 + self.action_dim],\n",
    "                all_other[:, 2 + self.action_dim:],\n",
    "                torch.as_tensor(self.buf_state[:self.now_len], device=self.device))\n",
    "\n",
    "    def update_now_len_before_sample(self):\n",
    "        \"\"\"update the a pointer `now_len`, which is the current data number of ReplayBuffer\n",
    "        \"\"\"\n",
    "        self.now_len = self.max_len if self.if_full else self.next_idx\n",
    "\n",
    "    def empty_buffer_before_explore(self):\n",
    "        \"\"\"we empty the buffer by set now_len=0. On-policy need to empty buffer before exploration\n",
    "        \"\"\"\n",
    "        self.next_idx = 0\n",
    "        self.now_len = 0\n",
    "        self.if_full = False\n",
    "\n",
    "    def print_state_norm(self, neg_avg=None, div_std=None):  # non-essential\n",
    "        \"\"\"print the state norm information: state_avg, state_std\n",
    "        We don't suggest to use running stat state.\n",
    "        We directly do normalization on state using the historical avg and std\n",
    "        eg. `state = (state + self.neg_state_avg) * self.div_state_std` in `PreprocessEnv.step_norm()`\n",
    "        neg_avg = -states.mean()\n",
    "        div_std = 1/(states.std()+1e-5) or 6/(states.max()-states.min())\n",
    "        :array neg_avg: neg_avg.shape=(state_dim)\n",
    "        :array div_std: div_std.shape=(state_dim)\n",
    "        \"\"\"\n",
    "        max_sample_size = 2 ** 14\n",
    "\n",
    "        '''check if pass'''\n",
    "        state_shape = self.buf_state.shape\n",
    "        if len(state_shape) > 2 or state_shape[1] > 64:\n",
    "            print(f\"| print_state_norm(): state_dim: {state_shape} is too large to print its norm. \")\n",
    "            return None\n",
    "\n",
    "        '''sample state'''\n",
    "        indices = np.arange(self.now_len)\n",
    "        rd.shuffle(indices)\n",
    "        indices = indices[:max_sample_size]  # len(indices) = min(self.now_len, max_sample_size)\n",
    "\n",
    "        batch_state = self.buf_state[indices]\n",
    "\n",
    "        '''compute state norm'''\n",
    "        if isinstance(batch_state, torch.Tensor):\n",
    "            batch_state = batch_state.cpu().data.numpy()\n",
    "        assert isinstance(batch_state, np.ndarray)\n",
    "\n",
    "        if batch_state.shape[1] > 64:\n",
    "            print(f\"| _print_norm(): state_dim: {batch_state.shape[1]:.0f} is too large to print its norm. \")\n",
    "            return None\n",
    "\n",
    "        if np.isnan(batch_state).any():  # 2020-12-12\n",
    "            batch_state = np.nan_to_num(batch_state)  # nan to 0\n",
    "\n",
    "        ary_avg = batch_state.mean(axis=0)\n",
    "        ary_std = batch_state.std(axis=0)\n",
    "        fix_std = ((np.max(batch_state, axis=0) - np.min(batch_state, axis=0)) / 6 + ary_std) / 2\n",
    "\n",
    "        if neg_avg is not None:  # norm transfer\n",
    "            ary_avg = ary_avg - neg_avg / div_std\n",
    "            ary_std = fix_std / div_std\n",
    "\n",
    "        print(f\"print_state_norm: state_avg, state_std (fixed)\")\n",
    "        print(f\"avg = np.{repr(ary_avg).replace('=float32', '=np.float32')}\")\n",
    "        print(f\"std = np.{repr(ary_std).replace('=float32', '=np.float32')}\")\n",
    "\n",
    "    def td_error_update(self, td_error):\n",
    "        self.tree.td_error_update(td_error)\n",
    "\n",
    "\n",
    "class ReplayBufferMP:\n",
    "    def __init__(self, max_len, state_dim, action_dim, if_on_policy, if_per, if_gpu, rollout_num):\n",
    "        \"\"\"Experience Replay Buffer for Multiple Processing\n",
    "        `int rollout_num` the rollout workers number\n",
    "        \"\"\"\n",
    "        self.now_len = 0\n",
    "        self.max_len = max_len\n",
    "        self.rollout_num = rollout_num\n",
    "\n",
    "        _max_len = max_len // rollout_num\n",
    "        self.buffers = [ReplayBuffer(_max_len, state_dim, action_dim, if_on_policy, if_per, if_gpu)\n",
    "                        for _ in range(rollout_num)]\n",
    "\n",
    "    def extend_buffer(self, state, other, i):\n",
    "        self.buffers[i].extend_buffer(state, other)\n",
    "\n",
    "    def sample_batch(self, batch_size) -> list:\n",
    "        # rd_batch_sizes = rd.rand(self.rollout_num)\n",
    "        # rd_batch_sizes = (rd_batch_sizes * (batch_size / rd_batch_sizes.sum())).astype(np.int)\n",
    "        # list_items = [self.buffers[i].sample_batch(rd_batch_sizes[i])\n",
    "        #               for i in range(self.rollout_num) if rd_batch_sizes[i] > 2]\n",
    "        bs = batch_size // self.rollout_num\n",
    "        list_items = [self.buffers[i].sample_batch(bs)\n",
    "                      for i in range(self.rollout_num)]\n",
    "\n",
    "        # list_items of reward, mask, action, state, next_state\n",
    "        # list_items of reward, mask, action, state, next_state, is_weights (PER)\n",
    "        return [torch.cat([item[i] for item in list_items], dim=0)\n",
    "                for i in range(len(list_items[0]))]\n",
    "\n",
    "    def sample_all(self) -> list:\n",
    "        l__r_m_a_n_s = [self.buffers[i].sample_all()\n",
    "                        for i in range(self.rollout_num)]\n",
    "        # list of reward, mask, action, noise, state\n",
    "        return [torch.cat([item[i] for item in l__r_m_a_n_s], dim=0)\n",
    "                for i in range(len(l__r_m_a_n_s[0]))]\n",
    "\n",
    "    def update_now_len_before_sample(self):\n",
    "        self.now_len = 0\n",
    "        for buffer in self.buffers:\n",
    "            buffer.update_now_len_before_sample()\n",
    "            self.now_len += buffer.now_len\n",
    "\n",
    "    def empty_buffer_before_explore(self):\n",
    "        for buffer in self.buffers:\n",
    "            buffer.empty_buffer_before_explore()\n",
    "\n",
    "    def print_state_norm(self, neg_avg=None, div_std=None):  # non-essential\n",
    "        # for buffer in self.l_buffer:\n",
    "        self.buffers[0].print_state_norm(neg_avg, div_std)\n",
    "\n",
    "    def td_error_update(self, td_error):\n",
    "        td_errors = td_error.view(self.rollout_num, -1, 1)\n",
    "        for i in range(self.rollout_num):\n",
    "            self.buffers[i].tree.td_error_update(td_errors[i])\n",
    "\n",
    "\n",
    "class BinarySearchTree:\n",
    "    \"\"\"Binary Search Tree for PER\n",
    "    Contributor: Github GyChou, Github mississippiu\n",
    "    Reference: https://github.com/kaixindelele/DRLib/tree/main/algos/pytorch/td3_sp\n",
    "    Reference: https://github.com/jaromiru/AI-blog/blob/master/SumTree.py\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, memo_len):\n",
    "        self.memo_len = memo_len  # replay buffer len\n",
    "        self.prob_ary = np.zeros((memo_len - 1) + memo_len)  # parent_nodes_num + leaf_nodes_num\n",
    "        self.max_len = len(self.prob_ary)\n",
    "        self.now_len = self.memo_len - 1  # pointer\n",
    "        self.indices = None\n",
    "        self.depth = int(np.log2(self.max_len))\n",
    "\n",
    "        # PER.  Prioritized Experience Replay. Section 4\n",
    "        # alpha, beta = 0.7, 0.5 for rank-based variant\n",
    "        # alpha, beta = 0.6, 0.4 for proportional variant\n",
    "        self.per_alpha = 0.6  # alpha = (Uniform:0, Greedy:1)\n",
    "        self.per_beta = 0.4  # beta = (PER:0, NotPER:1)\n",
    "\n",
    "    def update_id(self, data_id, prob=10):  # 10 is max_prob\n",
    "        tree_id = data_id + self.memo_len - 1\n",
    "        if self.now_len == tree_id:\n",
    "            self.now_len += 1\n",
    "\n",
    "        delta = prob - self.prob_ary[tree_id]\n",
    "        self.prob_ary[tree_id] = prob\n",
    "\n",
    "        while tree_id != 0:  # propagate the change through tree\n",
    "            tree_id = (tree_id - 1) // 2  # faster than the recursive loop\n",
    "            self.prob_ary[tree_id] += delta\n",
    "\n",
    "    def update_ids(self, data_ids, prob=10):  # 10 is max_prob\n",
    "        ids = data_ids + self.memo_len - 1\n",
    "        self.now_len += (ids >= self.now_len).sum()\n",
    "\n",
    "        upper_step = self.depth - 1\n",
    "        self.prob_ary[ids] = prob  # here, ids means the indices of given children (maybe the right ones or left ones)\n",
    "        p_ids = (ids - 1) // 2\n",
    "\n",
    "        while upper_step:  # propagate the change through tree\n",
    "            ids = p_ids * 2 + 1  # in this while loop, ids means the indices of the left children\n",
    "            self.prob_ary[p_ids] = self.prob_ary[ids] + self.prob_ary[ids + 1]\n",
    "            p_ids = (p_ids - 1) // 2\n",
    "            upper_step -= 1\n",
    "\n",
    "        self.prob_ary[0] = self.prob_ary[1] + self.prob_ary[2]\n",
    "        # because we take depth-1 upper steps, ps_tree[0] need to be updated alone\n",
    "\n",
    "    def get_leaf_id(self, v):\n",
    "        \"\"\"Tree structure and array storage:\n",
    "        Tree index:\n",
    "              0       -> storing priority sum\n",
    "            |  |\n",
    "          1     2\n",
    "         | |   | |\n",
    "        3  4  5  6    -> storing priority for transitions\n",
    "        Array type for storing: [0, 1, 2, 3, 4, 5, 6]\n",
    "        \"\"\"\n",
    "        parent_idx = 0\n",
    "        while True:\n",
    "            l_idx = 2 * parent_idx + 1  # the leaf's left node\n",
    "            r_idx = l_idx + 1  # the leaf's right node\n",
    "            if l_idx >= (len(self.prob_ary)):  # reach bottom, end search\n",
    "                leaf_idx = parent_idx\n",
    "                break\n",
    "            else:  # downward search, always search for a higher priority node\n",
    "                if v <= self.prob_ary[l_idx]:\n",
    "                    parent_idx = l_idx\n",
    "                else:\n",
    "                    v -= self.prob_ary[l_idx]\n",
    "                    parent_idx = r_idx\n",
    "        return min(leaf_idx, self.now_len - 2)  # leaf_idx\n",
    "\n",
    "    def get_indices_is_weights(self, batch_size, beg, end):\n",
    "        self.per_beta = min(1., self.per_beta + 0.001)\n",
    "\n",
    "        # get random values for searching indices with proportional prioritization\n",
    "        values = (rd.rand(batch_size) + np.arange(batch_size)) * (self.prob_ary[0] / batch_size)\n",
    "\n",
    "        # get proportional prioritization\n",
    "        leaf_ids = np.array([self.get_leaf_id(v) for v in values])\n",
    "        self.indices = leaf_ids - (self.memo_len - 1)\n",
    "\n",
    "        prob_ary = self.prob_ary[leaf_ids] / self.prob_ary[beg:end].min()\n",
    "        is_weights = np.power(prob_ary, -self.per_beta)  # important sampling weights\n",
    "        return self.indices, is_weights\n",
    "\n",
    "    def td_error_update(self, td_error):  # td_error = (q-q).detach_().abs()\n",
    "        prob = td_error.squeeze().clamp(1e-6, 10).pow(self.per_alpha)\n",
    "        prob = prob.cpu().numpy()\n",
    "        self.update_ids(self.indices, prob)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import gym\n",
    "import time\n",
    "import torch\n",
    "import numpy as np\n",
    "import numpy.random as rd\n",
    "from copy import deepcopy\n",
    "\n",
    "\"\"\"[ElegantRL](https://github.com/AI4Finance-LLC/ElegantRL)\"\"\"\n",
    "\n",
    "\n",
    "class Arguments:\n",
    "    def __init__(self, agent=None, env=None, gpu_id=None, if_on_policy=False):\n",
    "        self.agent = agent  # Deep Reinforcement Learning algorithm\n",
    "        self.cwd = None  # current work directory. cwd is None means set it automatically\n",
    "        self.env = env  # the environment for training\n",
    "        self.env_eval = None  # the environment for evaluating\n",
    "        self.gpu_id = gpu_id  # choose the GPU for running. gpu_id is None means set it automatically\n",
    "\n",
    "        '''Arguments for training (off-policy)'''\n",
    "        self.net_dim = 2 ** 8  # the network width\n",
    "        self.batch_size = 2 ** 8  # num of transitions sampled from replay buffer.\n",
    "        self.repeat_times = 2 ** 0  # repeatedly update network to keep critic's loss small\n",
    "        self.target_step = 2 ** 10  # collect target_step, then update network\n",
    "        self.max_memo = 2 ** 17  # capacity of replay buffer\n",
    "        if if_on_policy:  # (on-policy)\n",
    "            self.net_dim = 2 ** 9\n",
    "            self.batch_size = 2 ** 9\n",
    "            self.repeat_times = 2 ** 4\n",
    "            self.target_step = 2 ** 12\n",
    "            self.max_memo = self.target_step\n",
    "        self.gamma = 0.99  # discount factor of future rewards\n",
    "        self.reward_scale = 2 ** 0  # an approximate target reward usually be closed to 256\n",
    "        self.if_per = False  # Prioritized Experience Replay for sparse reward\n",
    "\n",
    "        self.rollout_num = 2  # the number of rollout workers (larger is not always faster)\n",
    "        self.num_threads = 8  # cpu_num for evaluate model, torch.set_num_threads(self.num_threads)\n",
    "\n",
    "        '''Arguments for evaluate'''\n",
    "        self.break_step = 2 ** 20  # break training after 'total_step > break_step'\n",
    "        self.if_remove = True  # remove the cwd folder? (True, False, None:ask me)\n",
    "        self.if_allow_break = True  # allow break training when reach goal (early termination)\n",
    "        self.eval_gap = 2 ** 5  # evaluate the agent per eval_gap seconds\n",
    "        self.eval_times1 = 2 ** 2  # evaluation times\n",
    "        self.eval_times2 = 2 ** 4  # evaluation times if 'eval_reward > max_reward'\n",
    "        self.random_seed = 0  # initialize random seed in self.init_before_training()\n",
    "\n",
    "    def init_before_training(self, if_main=True):\n",
    "        if self.agent is None:\n",
    "            raise RuntimeError('\\n| Why agent=None? Assignment args.agent = AgentXXX please.')\n",
    "        if not hasattr(self.agent, 'init'):\n",
    "            raise RuntimeError('\\n| There should be agent=AgentXXX() instead of agent=AgentXXX')\n",
    "        if self.env is None:\n",
    "            raise RuntimeError('\\n| Why env=None? Assignment args.env = XxxEnv() please.')\n",
    "        if isinstance(self.env, str) or not hasattr(self.env, 'env_name'):\n",
    "            raise RuntimeError('\\n| What is env.env_name? use env=PreprocessEnv(env). It is a Wrapper.')\n",
    "\n",
    "        '''set gpu_id automatically'''\n",
    "        if self.gpu_id is None:  # set gpu_id automatically\n",
    "            import sys\n",
    "            self.gpu_id = sys.argv[-1][-4]\n",
    "        else:\n",
    "            self.gpu_id = str(self.gpu_id)\n",
    "        if not self.gpu_id.isdigit():  # set gpu_id as '0' in default\n",
    "            self.gpu_id = '0'\n",
    "\n",
    "        '''set cwd automatically'''\n",
    "        if self.cwd is None:\n",
    "            agent_name = self.agent.__class__.__name__\n",
    "            self.cwd = f'./{agent_name}/{self.env.env_name}_{self.gpu_id}'\n",
    "\n",
    "        if if_main:\n",
    "            print(f'| GPU id: {self.gpu_id}, cwd: {self.cwd}')\n",
    "\n",
    "            import shutil  # remove history according to bool(if_remove)\n",
    "            if self.if_remove is None:\n",
    "                self.if_remove = bool(input(\"PRESS 'y' to REMOVE: {}? \".format(self.cwd)) == 'y')\n",
    "            if self.if_remove:\n",
    "                shutil.rmtree(self.cwd, ignore_errors=True)\n",
    "                print(\"| Remove history\")\n",
    "            os.makedirs(self.cwd, exist_ok=True)\n",
    "\n",
    "        os.environ['CUDA_VISIBLE_DEVICES'] = str(self.gpu_id)\n",
    "        torch.set_num_threads(self.num_threads)\n",
    "        torch.set_default_dtype(torch.float32)\n",
    "        torch.manual_seed(self.random_seed)\n",
    "        np.random.seed(self.random_seed)\n",
    "\n",
    "\n",
    "'''single process training'''\n",
    "\n",
    "\n",
    "def train_and_evaluate(args):\n",
    "    args.init_before_training()\n",
    "\n",
    "    '''basic arguments'''\n",
    "    cwd = args.cwd\n",
    "    env = args.env\n",
    "    agent = args.agent\n",
    "    gpu_id = args.gpu_id  # necessary for Evaluator?\n",
    "\n",
    "    '''training arguments'''\n",
    "    net_dim = args.net_dim\n",
    "    max_memo = args.max_memo\n",
    "    break_step = args.break_step\n",
    "    batch_size = args.batch_size\n",
    "    target_step = args.target_step\n",
    "    repeat_times = args.repeat_times\n",
    "    if_break_early = args.if_allow_break\n",
    "    if_per = args.if_per\n",
    "    gamma = args.gamma\n",
    "    reward_scale = args.reward_scale\n",
    "\n",
    "    '''evaluating arguments'''\n",
    "    eval_gap = args.eval_gap\n",
    "    eval_times1 = args.eval_times1\n",
    "    eval_times2 = args.eval_times2\n",
    "    if args.env_eval is not None:\n",
    "        env_eval = args.env_eval\n",
    "    elif args.env_eval in set(gym.envs.registry.env_specs.keys()):\n",
    "        env_eval = PreprocessEnv(gym.make(env.env_name))\n",
    "    else:\n",
    "        env_eval = deepcopy(env)\n",
    "\n",
    "    del args  # In order to show these hyper-parameters clearly, I put them above.\n",
    "\n",
    "    '''init: environment'''\n",
    "    max_step = env.max_step\n",
    "    state_dim = env.state_dim\n",
    "    action_dim = env.action_dim\n",
    "    if_discrete = env.if_discrete\n",
    "\n",
    "    '''init: Agent, ReplayBuffer, Evaluator'''\n",
    "    agent.init(net_dim, state_dim, action_dim, if_per)\n",
    "    if_on_policy = getattr(agent, 'if_on_policy', False)\n",
    "\n",
    "    buffer = ReplayBuffer(max_len=max_memo + max_step, state_dim=state_dim, action_dim=1 if if_discrete else action_dim,\n",
    "                          if_on_policy=if_on_policy, if_per=if_per, if_gpu=True)\n",
    "\n",
    "    evaluator = Evaluator(cwd=cwd, agent_id=gpu_id, device=agent.device, env=env_eval,\n",
    "                          eval_gap=eval_gap, eval_times1=eval_times1, eval_times2=eval_times2, )\n",
    "\n",
    "    '''prepare for training'''\n",
    "    agent.state = env.reset()\n",
    "    if if_on_policy:\n",
    "        steps = 0\n",
    "    else:  # explore_before_training for off-policy\n",
    "        with torch.no_grad():  # update replay buffer\n",
    "            steps = explore_before_training(env, buffer, target_step, reward_scale, gamma)\n",
    "\n",
    "        agent.update_net(buffer, target_step, batch_size, repeat_times)  # pre-training and hard update\n",
    "        agent.act_target.load_state_dict(agent.act.state_dict()) if getattr(agent, 'act_target', None) else None\n",
    "        agent.cri_target.load_state_dict(agent.cri.state_dict()) if getattr(agent, 'cri_target', None) else None\n",
    "    total_step = steps\n",
    "\n",
    "    '''start training'''\n",
    "    if_reach_goal = False\n",
    "    while not ((if_break_early and if_reach_goal)\n",
    "               or total_step > break_step\n",
    "               or os.path.exists(f'{cwd}/stop')):\n",
    "        steps = agent.explore_env(env, buffer, target_step, reward_scale, gamma)\n",
    "        total_step += steps\n",
    "\n",
    "        obj_a, obj_c = agent.update_net(buffer, target_step, batch_size, repeat_times)\n",
    "\n",
    "        if_reach_goal = evaluator.evaluate_save(agent.act, steps, obj_a, obj_c)\n",
    "        evaluator.draw_plot()\n",
    "\n",
    "    print(f'| SavedDir: {cwd}\\n| UsedTime: {time.time() - evaluator.start_time:.0f}')\n",
    "\n",
    "\n",
    "'''multiprocessing training'''\n",
    "\n",
    "\n",
    "def train_and_evaluate_mp(args):\n",
    "    act_workers = args.rollout_num\n",
    "    import multiprocessing as mp  # Python built-in multiprocessing library\n",
    "\n",
    "    pipe1_eva, pipe2_eva = mp.Pipe()  # Pipe() for Process mp_evaluate_agent()\n",
    "    pipe2_exp_list = list()  # Pipe() for Process mp_explore_in_env()\n",
    "\n",
    "    process_train = mp.Process(target=mp_train, args=(args, pipe2_eva, pipe2_exp_list))\n",
    "    process_evaluate = mp.Process(target=mp_evaluate, args=(args, pipe1_eva))\n",
    "    process = [process_train, process_evaluate]\n",
    "\n",
    "    for worker_id in range(act_workers):\n",
    "        exp_pipe1, exp_pipe2 = mp.Pipe(duplex=True)\n",
    "        pipe2_exp_list.append(exp_pipe1)\n",
    "        process.append(mp.Process(target=mp_explore, args=(args, exp_pipe2, worker_id)))\n",
    "\n",
    "    [p.start() for p in process]\n",
    "    process_evaluate.join()\n",
    "    process_train.join()\n",
    "    [p.terminate() for p in process]\n",
    "\n",
    "\n",
    "def mp_train(args, pipe1_eva, pipe1_exp_list):\n",
    "    args.init_before_training(if_main=False)\n",
    "\n",
    "    '''basic arguments'''\n",
    "    env = args.env\n",
    "    cwd = args.cwd\n",
    "    agent = args.agent\n",
    "    rollout_num = args.rollout_num\n",
    "\n",
    "    '''training arguments'''\n",
    "    net_dim = args.net_dim\n",
    "    max_memo = args.max_memo\n",
    "    break_step = args.break_step\n",
    "    batch_size = args.batch_size\n",
    "    target_step = args.target_step\n",
    "    repeat_times = args.repeat_times\n",
    "    if_break_early = args.if_allow_break\n",
    "    if_per = args.if_per\n",
    "    del args  # In order to show these hyper-parameters clearly, I put them above.\n",
    "\n",
    "    '''init: environment'''\n",
    "    max_step = env.max_step\n",
    "    state_dim = env.state_dim\n",
    "    action_dim = env.action_dim\n",
    "    if_discrete = env.if_discrete\n",
    "\n",
    "    '''init: Agent, ReplayBuffer'''\n",
    "    agent.init(net_dim, state_dim, action_dim, if_per)\n",
    "    if_on_policy = getattr(agent, 'if_on_policy', False)\n",
    "\n",
    "    '''send'''\n",
    "    pipe1_eva.send(agent.act)  # send\n",
    "    # act = pipe2_eva.recv()  # recv\n",
    "\n",
    "    buffer_mp = ReplayBufferMP(max_len=max_memo + max_step * rollout_num, if_on_policy=if_on_policy,\n",
    "                               state_dim=state_dim, action_dim=1 if if_discrete else action_dim,\n",
    "                               rollout_num=rollout_num, if_gpu=True, if_per=if_per)\n",
    "\n",
    "    '''prepare for training'''\n",
    "    if if_on_policy:\n",
    "        steps = 0\n",
    "    else:  # explore_before_training for off-policy\n",
    "        with torch.no_grad():  # update replay buffer\n",
    "            steps = 0\n",
    "            for i in range(rollout_num):\n",
    "                pipe1_exp = pipe1_exp_list[i]\n",
    "\n",
    "                # pipe2_exp.send((buffer.buf_state[:buffer.now_len], buffer.buf_other[:buffer.now_len]))\n",
    "                buf_state, buf_other = pipe1_exp.recv()\n",
    "\n",
    "                steps += len(buf_state)\n",
    "                buffer_mp.extend_buffer(buf_state, buf_other, i)\n",
    "\n",
    "        agent.update_net(buffer_mp, target_step, batch_size, repeat_times)  # pre-training and hard update\n",
    "        agent.act_target.load_state_dict(agent.act.state_dict()) if getattr(env, 'act_target', None) else None\n",
    "        agent.cri_target.load_state_dict(agent.cri.state_dict()) if getattr(env, 'cri_target', None) in dir(\n",
    "            agent) else None\n",
    "    total_step = steps\n",
    "    '''send'''\n",
    "    pipe1_eva.send((agent.act, steps, 0, 0.5))  # send\n",
    "    # act, steps, obj_a, obj_c = pipe2_eva.recv()  # recv\n",
    "\n",
    "    '''start training'''\n",
    "    if_solve = False\n",
    "    while not ((if_break_early and if_solve)\n",
    "               or total_step > break_step\n",
    "               or os.path.exists(f'{cwd}/stop')):\n",
    "        '''update ReplayBuffer'''\n",
    "        steps = 0  # send by pipe1_eva\n",
    "        for i in range(rollout_num):\n",
    "            pipe1_exp = pipe1_exp_list[i]\n",
    "            '''send'''\n",
    "            pipe1_exp.send(agent.act)\n",
    "            # agent.act = pipe2_exp.recv()\n",
    "            '''recv'''\n",
    "            # pipe2_exp.send((buffer.buf_state[:buffer.now_len], buffer.buf_other[:buffer.now_len]))\n",
    "            buf_state, buf_other = pipe1_exp.recv()\n",
    "\n",
    "            steps += len(buf_state)\n",
    "            buffer_mp.extend_buffer(buf_state, buf_other, i)\n",
    "        total_step += steps\n",
    "\n",
    "        '''update network parameters'''\n",
    "        obj_a, obj_c = agent.update_net(buffer_mp, target_step, batch_size, repeat_times)\n",
    "\n",
    "        '''saves the agent with max reward'''\n",
    "        '''send'''\n",
    "        pipe1_eva.send((agent.act, steps, obj_a, obj_c))\n",
    "        # q_i_eva_get = pipe2_eva.recv()\n",
    "\n",
    "        if_solve = pipe1_eva.recv()\n",
    "\n",
    "        if pipe1_eva.poll():\n",
    "            '''recv'''\n",
    "            # pipe2_eva.send(if_solve)\n",
    "            if_solve = pipe1_eva.recv()\n",
    "\n",
    "    buffer_mp.print_state_norm(env.neg_state_avg if hasattr(env, 'neg_state_avg') else None,\n",
    "                               env.div_state_std if hasattr(env, 'div_state_std') else None)  # 2020-12-12\n",
    "\n",
    "    '''send'''\n",
    "    pipe1_eva.send('stop')\n",
    "    # q_i_eva_get = pipe2_eva.recv()\n",
    "    time.sleep(4)\n",
    "\n",
    "\n",
    "def mp_explore(args, pipe2_exp, worker_id):\n",
    "    args.init_before_training(if_main=False)\n",
    "\n",
    "    '''basic arguments'''\n",
    "    env = args.env\n",
    "    agent = args.agent\n",
    "    rollout_num = args.rollout_num\n",
    "\n",
    "    '''training arguments'''\n",
    "    net_dim = args.net_dim\n",
    "    max_memo = args.max_memo\n",
    "    target_step = args.target_step\n",
    "    gamma = args.gamma\n",
    "    if_per = args.if_per\n",
    "    reward_scale = args.reward_scale\n",
    "\n",
    "    random_seed = args.random_seed\n",
    "    torch.manual_seed(random_seed + worker_id)\n",
    "    np.random.seed(random_seed + worker_id)\n",
    "    del args  # In order to show these hyper-parameters clearly, I put them above.\n",
    "\n",
    "    '''init: environment'''\n",
    "    max_step = env.max_step\n",
    "    state_dim = env.state_dim\n",
    "    action_dim = env.action_dim\n",
    "    if_discrete = env.if_discrete\n",
    "\n",
    "    '''init: Agent, ReplayBuffer'''\n",
    "    agent.init(net_dim, state_dim, action_dim, if_per)\n",
    "    agent.state = env.reset()\n",
    "\n",
    "    if_on_policy = getattr(agent, 'if_on_policy', False)\n",
    "    buffer = ReplayBuffer(max_len=max_memo // rollout_num + max_step, if_on_policy=if_on_policy,\n",
    "                          state_dim=state_dim, action_dim=1 if if_discrete else action_dim,\n",
    "                          if_per=if_per, if_gpu=False)\n",
    "\n",
    "    '''start exploring'''\n",
    "    exp_step = target_step // rollout_num\n",
    "    with torch.no_grad():\n",
    "        if not if_on_policy:\n",
    "            explore_before_training(env, buffer, exp_step, reward_scale, gamma)\n",
    "\n",
    "            buffer.update_now_len_before_sample()\n",
    "\n",
    "            pipe2_exp.send((buffer.buf_state[:buffer.now_len], buffer.buf_other[:buffer.now_len]))\n",
    "            # buf_state, buf_other = pipe1_exp.recv()\n",
    "\n",
    "            buffer.empty_buffer_before_explore()\n",
    "\n",
    "        while True:\n",
    "            agent.explore_env(env, buffer, exp_step, reward_scale, gamma)\n",
    "\n",
    "            buffer.update_now_len_before_sample()\n",
    "\n",
    "            pipe2_exp.send((buffer.buf_state[:buffer.now_len], buffer.buf_other[:buffer.now_len]))\n",
    "            # buf_state, buf_other = pipe1_exp.recv()\n",
    "\n",
    "            buffer.empty_buffer_before_explore()\n",
    "\n",
    "            # pipe1_exp.send(agent.act)\n",
    "            agent.act = pipe2_exp.recv()\n",
    "\n",
    "\n",
    "def mp_evaluate(args, pipe2_eva):\n",
    "    args.init_before_training(if_main=True)\n",
    "\n",
    "    '''basic arguments'''\n",
    "    cwd = args.cwd\n",
    "    env = args.env\n",
    "    env_eval = env if args.env_eval is None else args.env_eval\n",
    "    agent_id = args.gpu_id\n",
    "\n",
    "    '''evaluating arguments'''\n",
    "    eval_gap = args.eval_gap\n",
    "    eval_times1 = args.eval_times1\n",
    "    eval_times2 = args.eval_times2\n",
    "    del args  # In order to show these hyper-parameters clearly, I put them above.\n",
    "\n",
    "    '''init: Evaluator'''\n",
    "    evaluator = Evaluator(cwd=cwd, agent_id=agent_id, device=torch.device(\"cpu\"), env=env_eval,\n",
    "                          eval_gap=eval_gap, eval_times1=eval_times1, eval_times2=eval_times2, )  # build Evaluator\n",
    "\n",
    "    '''act_cpu without gradient for pipe1_eva'''\n",
    "    # pipe1_eva.send(agent.act)\n",
    "    act = pipe2_eva.recv()\n",
    "\n",
    "    act_cpu = deepcopy(act).to(torch.device(\"cpu\"))  # for pipe1_eva\n",
    "    [setattr(param, 'requires_grad', False) for param in act_cpu.parameters()]\n",
    "\n",
    "    '''start evaluating'''\n",
    "    with torch.no_grad():  # speed up running\n",
    "        act, steps, obj_a, obj_c = pipe2_eva.recv()  # pipe2_eva (act, steps, obj_a, obj_c)\n",
    "\n",
    "        if_loop = True\n",
    "        while if_loop:\n",
    "            '''update actor'''\n",
    "            while not pipe2_eva.poll():  # wait until pipe2_eva not empty\n",
    "                time.sleep(1)\n",
    "            steps_sum = 0\n",
    "            while pipe2_eva.poll():  # receive the latest object from pipe\n",
    "                '''recv'''\n",
    "                # pipe1_eva.send((agent.act, steps, obj_a, obj_c))\n",
    "                # pipe1_eva.send('stop')\n",
    "                q_i_eva_get = pipe2_eva.recv()\n",
    "\n",
    "                if q_i_eva_get == 'stop':\n",
    "                    if_loop = False\n",
    "                    break\n",
    "                act, steps, obj_a, obj_c = q_i_eva_get\n",
    "                steps_sum += steps\n",
    "            act_cpu.load_state_dict(act.state_dict())\n",
    "            if_solve = evaluator.evaluate_save(act_cpu, steps_sum, obj_a, obj_c)\n",
    "            '''send'''\n",
    "            pipe2_eva.send(if_solve)\n",
    "            # if_solve = pipe1_eva.recv()\n",
    "\n",
    "            evaluator.draw_plot()\n",
    "\n",
    "    print(f'| SavedDir: {cwd}\\n| UsedTime: {time.time() - evaluator.start_time:.0f}')\n",
    "\n",
    "    while pipe2_eva.poll():  # empty the pipe\n",
    "        pipe2_eva.recv()\n",
    "\n",
    "\n",
    "'''utils'''\n",
    "\n",
    "\n",
    "class Evaluator:\n",
    "    def __init__(self, cwd, agent_id, eval_times1, eval_times2, eval_gap, env, device):\n",
    "        self.recorder = [(0., -np.inf, 0., 0., 0.), ]  # total_step, r_avg, r_std, obj_a, obj_c\n",
    "        self.r_max = -np.inf\n",
    "        self.total_step = 0\n",
    "\n",
    "        self.cwd = cwd  # constant\n",
    "        self.device = device\n",
    "        self.agent_id = agent_id\n",
    "        self.eval_gap = eval_gap\n",
    "        self.eval_times1 = eval_times1\n",
    "        self.eval_times2 = eval_times2\n",
    "        self.env = env\n",
    "        self.target_return = env.target_return\n",
    "\n",
    "        self.used_time = None\n",
    "        self.start_time = time.time()\n",
    "        self.eval_time = -1  # a early time\n",
    "        print(f\"{'ID':>2}  {'Step':>8}  {'MaxR':>8} |\"\n",
    "              f\"{'avgR':>8}  {'stdR':>8}   {'objA':>8}  {'objC':>8} |\"\n",
    "              f\"{'avgS':>6}  {'stdS':>4}\")\n",
    "\n",
    "    def evaluate_save(self, act, steps, obj_a, obj_c) -> bool:\n",
    "        self.total_step += steps  # update total training steps\n",
    "\n",
    "        if time.time() - self.eval_time > self.eval_gap:\n",
    "            self.eval_time = time.time()\n",
    "\n",
    "            rewards_steps_list = [get_episode_return(self.env, act, self.device) for _ in range(self.eval_times1)]\n",
    "            r_avg, r_std, s_avg, s_std = self.get_r_avg_std_s_avg_std(rewards_steps_list)\n",
    "\n",
    "            if r_avg > self.r_max:  # evaluate actor twice to save CPU Usage and keep precision\n",
    "                rewards_steps_list += [get_episode_return(self.env, act, self.device)\n",
    "                                       for _ in range(self.eval_times2 - self.eval_times1)]\n",
    "                r_avg, r_std, s_avg, s_std = self.get_r_avg_std_s_avg_std(rewards_steps_list)\n",
    "            if r_avg > self.r_max:  # save checkpoint with highest episode return\n",
    "                self.r_max = r_avg  # update max reward (episode return)\n",
    "\n",
    "                '''save actor.pth'''\n",
    "                act_save_path = f'{self.cwd}/actor.pth'\n",
    "                torch.save(act.state_dict(), act_save_path)\n",
    "                print(f\"{self.agent_id:<2}  {self.total_step:8.2e}  {self.r_max:8.2f} |\")  # save policy and print\n",
    "\n",
    "            self.recorder.append((self.total_step, r_avg, r_std, obj_a, obj_c))  # update recorder\n",
    "\n",
    "            if_reach_goal = bool(self.r_max > self.target_return)  # check if_reach_goal\n",
    "            if if_reach_goal and self.used_time is None:\n",
    "                self.used_time = int(time.time() - self.start_time)\n",
    "                print(f\"{'ID':>2}  {'Step':>8}  {'TargetR':>8} |\"\n",
    "                      f\"{'avgR':>8}  {'stdR':>8}   {'UsedTime':>8}  ########\\n\"\n",
    "                      f\"{self.agent_id:<2}  {self.total_step:8.2e}  {self.target_return:8.2f} |\"\n",
    "                      f\"{r_avg:8.2f}  {r_std:8.2f}   {self.used_time:>8}  ########\")\n",
    "\n",
    "            print(f\"{self.agent_id:<2}  {self.total_step:8.2e}  {self.r_max:8.2f} |\"\n",
    "                  f\"{r_avg:8.2f}  {r_std:8.2f}   {obj_a:8.2f}  {obj_c:8.2f} |\"\n",
    "                  f\"{s_avg:6.0f}  {s_std:4.0f}\")\n",
    "        else:\n",
    "            if_reach_goal = False\n",
    "        return if_reach_goal\n",
    "\n",
    "    def draw_plot(self):\n",
    "        if len(self.recorder) == 0:\n",
    "            print(\"| save_npy_draw_plot() WARNNING: len(self.recorder)==0\")\n",
    "            return None\n",
    "\n",
    "        '''convert to array and save as npy'''\n",
    "        np.save('%s/recorder.npy' % self.cwd, self.recorder)\n",
    "\n",
    "        '''draw plot and save as png'''\n",
    "        train_time = int(time.time() - self.start_time)\n",
    "        total_step = int(self.recorder[-1][0])\n",
    "        save_title = f\"plot_step_time_maxR_{int(total_step)}_{int(train_time)}_{self.r_max:.3f}\"\n",
    "\n",
    "        save_learning_curve(self.recorder, self.cwd, save_title)\n",
    "\n",
    "    @staticmethod\n",
    "    def get_r_avg_std_s_avg_std(rewards_steps_list):\n",
    "        rewards_steps_ary = np.array(rewards_steps_list)\n",
    "        r_avg, s_avg = rewards_steps_ary.mean(axis=0)  # average of episode return and episode step\n",
    "        r_std, s_std = rewards_steps_ary.std(axis=0)  # standard dev. of episode return and episode step\n",
    "        return r_avg, r_std, s_avg, s_std\n",
    "\n",
    "\n",
    "def get_episode_return(env, act, device) -> (float, int):\n",
    "    episode_return = 0.0  # sum of rewards in an episode\n",
    "    episode_step = 1\n",
    "    max_step = env.max_step\n",
    "    if_discrete = env.if_discrete\n",
    "\n",
    "    state = env.reset()\n",
    "    for episode_step in range(max_step):\n",
    "        s_tensor = torch.as_tensor((state,), device=device)\n",
    "        a_tensor = act(s_tensor)\n",
    "        if if_discrete:\n",
    "            a_tensor = a_tensor.argmax(dim=1)\n",
    "        action = a_tensor.detach().cpu().numpy()[0]  # not need detach(), because with torch.no_grad() outside\n",
    "        state, reward, done, _ = env.step(action)\n",
    "        episode_return += reward\n",
    "        if done:\n",
    "            break\n",
    "    episode_return = getattr(env, 'episode_return', episode_return)\n",
    "    return episode_return, episode_step + 1\n",
    "\n",
    "\n",
    "def save_learning_curve(recorder, cwd='.', save_title='learning curve'):\n",
    "    recorder = np.array(recorder)  # recorder_ary.append((self.total_step, r_avg, r_std, obj_a, obj_c))\n",
    "    steps = recorder[:, 0]  # x-axis is training steps\n",
    "    r_avg = recorder[:, 1]\n",
    "    r_std = recorder[:, 2]\n",
    "    obj_a = recorder[:, 3]\n",
    "    obj_c = recorder[:, 4]\n",
    "\n",
    "    '''plot subplots'''\n",
    "    import matplotlib as mpl\n",
    "    mpl.use('Agg')\n",
    "    \"\"\"Generating matplotlib graphs without a running X server [duplicate]\n",
    "    write `mpl.use('Agg')` before `import matplotlib.pyplot as plt`\n",
    "    https://stackoverflow.com/a/4935945/9293137\n",
    "    \"\"\"\n",
    "    import matplotlib.pyplot as plt\n",
    "    fig, axs = plt.subplots(2)\n",
    "\n",
    "    axs0 = axs[0]\n",
    "    axs0.cla()\n",
    "    color0 = 'lightcoral'\n",
    "    axs0.set_xlabel('Total Steps')\n",
    "    axs0.set_ylabel('Episode Return')\n",
    "    axs0.plot(steps, r_avg, label='Episode Return', color=color0)\n",
    "    axs0.fill_between(steps, r_avg - r_std, r_avg + r_std, facecolor=color0, alpha=0.3)\n",
    "\n",
    "    ax11 = axs[1]\n",
    "    ax11.cla()\n",
    "    color11 = 'royalblue'\n",
    "    axs0.set_xlabel('Total Steps')\n",
    "    ax11.set_ylabel('objA', color=color11)\n",
    "    ax11.plot(steps, obj_a, label='objA', color=color11)\n",
    "    ax11.tick_params(axis='y', labelcolor=color11)\n",
    "\n",
    "    ax12 = axs[1].twinx()\n",
    "    color12 = 'darkcyan'\n",
    "    ax12.set_ylabel('objC', color=color12)\n",
    "    ax12.fill_between(steps, obj_c, facecolor=color12, alpha=0.2, )\n",
    "    ax12.tick_params(axis='y', labelcolor=color12)\n",
    "\n",
    "    '''plot save'''\n",
    "    plt.title(save_title, y=2.3)\n",
    "    plt.savefig(f\"{cwd}/plot_learning_curve.jpg\")\n",
    "    plt.close('all')  # avoiding warning about too many open figures, rcParam `figure.max_open_warning`\n",
    "    # plt.show()  # if use `mpl.use('Agg')` to draw figures without GUI, then plt can't plt.show()\n",
    "\n",
    "\n",
    "def explore_before_training(env, buffer, target_step, reward_scale, gamma) -> int:\n",
    "    # just for off-policy. Because on-policy don't explore before training.\n",
    "    if_discrete = env.if_discrete\n",
    "    action_dim = env.action_dim\n",
    "\n",
    "    state = env.reset()\n",
    "    steps = 0\n",
    "\n",
    "    while steps < target_step:\n",
    "        action = rd.randint(action_dim) if if_discrete else rd.uniform(-1, 1, size=action_dim)\n",
    "        next_state, reward, done, _ = env.step(action)\n",
    "        steps += 1\n",
    "\n",
    "        scaled_reward = reward * reward_scale\n",
    "        mask = 0.0 if done else gamma\n",
    "        other = (scaled_reward, mask, action) if if_discrete else (scaled_reward, mask, *action)\n",
    "        buffer.append_buffer(state, other)\n",
    "\n",
    "        state = env.reset() if done else next_state\n",
    "    return steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1950, 15)\n",
      "| GPU id: 0, cwd: ./AgentPPO/Stock_qc-v1_0\n",
      "| Remove history\n",
      "ID      Step      MaxR |    avgR      stdR       objA      objC |  avgS  stdS\n",
      "0.9906872846391018\n"
     ]
    }
   ],
   "source": [
    "stock_list = [\n",
    "    \"AAPL\",\"MSFT\",\"JPM\"\n",
    "]\n",
    "start = datetime(2021,5,1)\n",
    "end = datetime(2021,5,10)\n",
    "resolution = Resolution.Minute\n",
    "df = data_fetch(start,end,stock_list,resolution)\n",
    "ary = preprocess(df, stock_list)\n",
    "print(ary.shape)\n",
    "args = Arguments(if_on_policy=True)\n",
    "args.agent = AgentPPO()\n",
    "\n",
    "'''choose environment'''\n",
    "args.env = StockTradingEnv(ary = ary, if_train=True)\n",
    "args.env_eval = StockTradingEnv(ary = ary, if_train=False)\n",
    "args.net_dim = 2 ** 9 # change a default hyper-parameters\n",
    "args.batch_size = 2 ** 8\n",
    "args.break_step = int(2e4)\n",
    "\n",
    "train_and_evaluate(args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
